From 055af12cf54726759b789c7941e280cb6572a9e4 Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:36:17 +0800 Subject: [PATCH 1/5] test(api): replace test doubles in KnowledgeFS extension and repository tests Co-Authored-By: Claude Opus 5.5 --- .../unit_tests/extensions/test_celery_ssl.py | 6 +- .../test_ext_knowledge_fs_observability.py | 18 ++-- .../extensions/test_knowledge_fs_celery.py | 92 +++++++++---------- ...y_knowledge_fs_control_space_repository.py | 6 +- ...nowledge_fs_lifecycle_outbox_repository.py | 12 +-- .../tasks/knowledge_fs_task_fakes.py | 42 +++++++++ .../unit_tests/test_celery_entrypoint.py | 5 +- 7 files changed, 113 insertions(+), 68 deletions(-) create mode 100644 api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py diff --git a/api/tests/unit_tests/extensions/test_celery_ssl.py b/api/tests/unit_tests/extensions/test_celery_ssl.py index 0614ae2f483..56eed648513 100644 --- a/api/tests/unit_tests/extensions/test_celery_ssl.py +++ b/api/tests/unit_tests/extensions/test_celery_ssl.py @@ -2,10 +2,10 @@ import ssl from importlib import import_module -from types import SimpleNamespace from unittest.mock import MagicMock, patch from enums import DeploymentEdition +from services.knowledge_fs.lifecycle_readiness import KnowledgeFSLifecycleWorkerReadiness class TestCelerySSLConfiguration: @@ -188,7 +188,7 @@ class TestCelerySSLConfiguration: patch("extensions.ext_celery.dify_config", mock_config), patch( "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=SimpleNamespace(ready=False), + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()), ), ): from dify_app import DifyApp @@ -249,7 +249,7 @@ class TestCelerySSLConfiguration: patch("extensions.ext_celery.dify_config", mock_config), patch( "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=SimpleNamespace(ready=False), + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()), ), ): from dify_app import DifyApp 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 29de6e93d06..7e8e5a661c2 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 @@ -1,17 +1,22 @@ from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import patch from flask import Flask +from sqlalchemy.orm import sessionmaker 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 def test_extension_registers_the_database_backed_control_space_gauge() -> None: app = Flask(__name__) - session_maker = MagicMock() - metrics = MagicMock() - reader = MagicMock() + session_maker = sessionmaker() + metrics = RecordingKnowledgeFSMetrics() + reader = SQLAlchemyKnowledgeFSControlSpaceStateCountReader(session_maker) with ( patch( @@ -30,13 +35,12 @@ def test_extension_registers_the_database_backed_control_space_gauge() -> None: ext_knowledge_fs_observability.init_app(app) reader_type.assert_called_once_with(session_maker) - metrics.register_control_space_state_gauge.assert_called_once_with(reader) + assert metrics.gauge_readers == [reader] def test_extension_does_not_block_startup_when_metric_registration_fails() -> None: app = Flask(__name__) - metrics = MagicMock() - metrics.register_control_space_state_gauge.side_effect = RuntimeError("metrics unavailable") + metrics = RecordingKnowledgeFSMetrics(gauge_registration_error=RuntimeError("metrics unavailable")) with ( patch("extensions.ext_knowledge_fs_observability.session_factory.get_session_maker"), diff --git a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py b/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py index d79a8f8740b..32a11e709b5 100644 --- a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py +++ b/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py @@ -1,12 +1,13 @@ from datetime import timedelta -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest 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.mark.parametrize("background_enabled", [False, True]) @@ -16,53 +17,51 @@ def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_rea ) -> None: # A stale deployment variable must not suppress vector cleanup for the selected provider. monkeypatch.setenv("KNOWLEDGE_VECTOR_STORAGE", "postgres") - config = MagicMock() - config.BROKER_USE_SSL = False - config.REDIS_KEY_PREFIX = "test" - config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1 - config.CELERY_BROKER_URL = "redis://localhost:6379/0" - config.CELERY_BACKEND = "redis" - config.CELERY_RESULT_BACKEND = "redis://localhost:6379/0" - config.CELERY_USE_SENTINEL = False - config.CELERY_BROKER_VISIBILITY_TIMEOUT = 10800 - config.KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED = background_enabled - config.VECTOR_STORE = vector_store - config.KNOWLEDGE_FS_BACKGROUND_POLL_INTERVAL_SECONDS = 5 - config.KNOWLEDGE_FS_BACKGROUND_TASK_TIMEOUT_SECONDS = 7200 - config.LOG_FORMAT = "%(message)s" - config.LOG_TZ = "UTC" - config.LOG_FILE = None - config.CELERY_TASK_ANNOTATIONS = dict[str, object]() - config.CELERY_BEAT_SCHEDULER_TIME = 1 - config.KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS = 2 - config.ENABLE_CONVERSATION_CLEANUP_TASK = True - config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 - config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False - config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False - config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False - config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK = False - config.ENABLE_CLEAN_MESSAGES = False - config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK = False - config.ENABLE_DATASETS_QUEUE_MONITOR = False - config.ENABLE_HUMAN_INPUT_TIMEOUT_TASK = False - config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK = False - config.MARKETPLACE_ENABLED = False - config.WORKFLOW_LOG_CLEANUP_ENABLED = False - config.ENABLE_WORKFLOW_RUN_CLEANUP_TASK = False - config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK = False - config.WORKFLOW_SCHEDULE_POLLER_INTERVAL = 1 - config.ENABLE_TRIGGER_PROVIDER_REFRESH_TASK = False - config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15 - config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False - config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30 - config.ENTERPRISE_ENABLED = False - config.ENTERPRISE_TELEMETRY_ENABLED = False + 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("extensions.ext_celery.dify_config", config), patch( "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=SimpleNamespace(ready=True), + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=True, blockers=()), ), ): celery_app = init_app(DifyApp(__name__)) @@ -99,10 +98,9 @@ def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_rea assert not any(name.startswith("knowledge_fs_background_") for name in schedules) with ( - patch("extensions.ext_celery.dify_config", config), patch( "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=SimpleNamespace(ready=False), + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=("test",)), ), ): preview_only_app = init_app(DifyApp(f"{__name__}.preview_only")) 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 fb779f1bb66..972869cbcdd 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 @@ -1,7 +1,6 @@ from __future__ import annotations from datetime import UTC, datetime -from unittest.mock import MagicMock import pytest from sqlalchemy.orm import Session @@ -18,6 +17,7 @@ from services.knowledge_fs.control_space_lifecycle import ( KnowledgeFSSpaceRegistrationConflictError, KnowledgeFSWorkspaceDeletionBlockedError, ) +from tests.unit_tests.tasks.knowledge_fs_task_fakes import RecordingKnowledgeFSMetrics @pytest.fixture @@ -54,7 +54,7 @@ def test_lifecycle_service_uses_compare_and_set_resource_versions( ) sqlite_session.add(control_space) sqlite_session.commit() - metrics = MagicMock() + metrics = RecordingKnowledgeFSMetrics() service = KnowledgeFSControlSpaceLifecycleService(repository, metrics=metrics) transitioned = service.transition( @@ -69,7 +69,7 @@ def test_lifecycle_service_uses_compare_and_set_resource_versions( assert transitioned.state is KnowledgeFSControlSpaceState.ACTIVE assert transitioned.resource_version == 1 assert transitioned.knowledge_space_id == "knowledge-space-1" - state_metric = metrics.record_control_space_state.call_args.args[0] + state_metric = metrics.control_space_states[-1] assert state_metric.from_state == "provisioning" assert state_metric.to_state == "active" assert control_space.id not in str(state_metric) 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 888e8aec7d1..f9c50bd8add 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 @@ -1,7 +1,6 @@ from __future__ import annotations from datetime import timedelta -from unittest.mock import MagicMock import pytest from sqlalchemy.orm import Session @@ -17,6 +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 def _outbox(control_space: KnowledgeFSControlSpace) -> KnowledgeFSLifecycleOutbox: @@ -53,15 +53,15 @@ def test_add_records_only_a_low_cardinality_queued_observation(sqlite_session: S ) sqlite_session.add(control_space) sqlite_session.flush() - metrics = MagicMock() + metrics = RecordingKnowledgeFSMetrics() repository = SQLAlchemyKnowledgeFSLifecycleOutboxRepository(sqlite_session, metrics=metrics) command = repository.add(_outbox(control_space)) - metrics.record_lifecycle_task.assert_called_once_with(KnowledgeFSLifecycleTaskMetric(None, "provision", "queued")) - assert command.id not in str(metrics.record_lifecycle_task.call_args_list) - assert control_space.id not in str(metrics.record_lifecycle_task.call_args_list) - assert control_space.tenant_id not in str(metrics.record_lifecycle_task.call_args_list) + assert metrics.lifecycle_tasks == [KnowledgeFSLifecycleTaskMetric(None, "provision", "queued")] + assert command.id not in str(metrics.lifecycle_tasks) + assert control_space.id not in str(metrics.lifecycle_tasks) + assert control_space.tenant_id not in str(metrics.lifecycle_tasks) @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py b/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py new file mode 100644 index 00000000000..86fcca07b64 --- /dev/null +++ b/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py @@ -0,0 +1,42 @@ +"""Typed fakes shared by the KnowledgeFS task, command and extension unit tests.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field + +from services.knowledge_fs.observability import ( + KnowledgeFSBatchStatusMetric, + KnowledgeFSCapabilityIssuanceMetric, + KnowledgeFSControlSpaceStateMetric, + KnowledgeFSLifecycleTaskMetric, +) + + +@dataclass +class RecordingKnowledgeFSMetrics: + """In-memory ``KnowledgeFSOperationalMetricsPort`` that records every event.""" + + batch_statuses: list[KnowledgeFSBatchStatusMetric] = field(default_factory=list) + capability_issuances: list[KnowledgeFSCapabilityIssuanceMetric] = field(default_factory=list) + control_space_states: list[KnowledgeFSControlSpaceStateMetric] = field(default_factory=list) + lifecycle_tasks: list[KnowledgeFSLifecycleTaskMetric] = field(default_factory=list) + gauge_readers: list[Callable[[], Mapping[str, int]]] = field(default_factory=list) + gauge_registration_error: Exception | None = None + + def record_batch_status(self, event: KnowledgeFSBatchStatusMetric) -> None: + self.batch_statuses.append(event) + + def record_capability_issuance(self, event: KnowledgeFSCapabilityIssuanceMetric) -> None: + self.capability_issuances.append(event) + + def record_control_space_state(self, event: KnowledgeFSControlSpaceStateMetric) -> None: + self.control_space_states.append(event) + + def record_lifecycle_task(self, event: KnowledgeFSLifecycleTaskMetric) -> None: + self.lifecycle_tasks.append(event) + + def register_control_space_state_gauge(self, read_counts: Callable[[], Mapping[str, int]]) -> None: + if self.gauge_registration_error is not None: + raise self.gauge_registration_error + self.gauge_readers.append(read_counts) diff --git a/api/tests/unit_tests/test_celery_entrypoint.py b/api/tests/unit_tests/test_celery_entrypoint.py index b7dcd7acca9..d85ab341224 100644 --- a/api/tests/unit_tests/test_celery_entrypoint.py +++ b/api/tests/unit_tests/test_celery_entrypoint.py @@ -1,7 +1,7 @@ import runpy import sys from pathlib import Path -from types import SimpleNamespace +from types import ModuleType from unittest.mock import patch import pytest @@ -9,7 +9,8 @@ import pytest @pytest.mark.parametrize("gevent_pool", [False, True]) def test_entrypoint_only_patches_libraries_for_a_gevent_pool(gevent_pool: bool) -> None: - fake_app = SimpleNamespace(app=object(), celery=object()) + fake_app = ModuleType("app") + fake_app.__dict__.update(app=object(), celery=object()) with ( patch.dict(sys.modules, {"app": fake_app}), patch("gevent.monkey.is_module_patched", return_value=gevent_pool), From e3294aef885a340d949a94d5daf285b587883f50 Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:38:34 +0800 Subject: [PATCH 2/5] test(api): use real session and recording bindings in KnowledgeFS publish handler test Co-Authored-By: Claude Opus 5.5 --- ...ngs_when_app_published_workflow_updated.py | 100 ++++++++++++------ 1 file changed, 66 insertions(+), 34 deletions(-) diff --git a/api/tests/unit_tests/events/event_handlers/test_sync_knowledge_fs_bindings_when_app_published_workflow_updated.py b/api/tests/unit_tests/events/event_handlers/test_sync_knowledge_fs_bindings_when_app_published_workflow_updated.py index 512e1e2d0ea..0752098b1ee 100644 --- a/api/tests/unit_tests/events/event_handlers/test_sync_knowledge_fs_bindings_when_app_published_workflow_updated.py +++ b/api/tests/unit_tests/events/event_handlers/test_sync_knowledge_fs_bindings_when_app_published_workflow_updated.py @@ -1,17 +1,47 @@ from __future__ import annotations import json -from types import SimpleNamespace -from unittest.mock import Mock +from collections.abc import Sequence +from dataclasses import dataclass, field import pytest from sqlalchemy.orm import Session from core.workflow.nodes.knowledge_retrieval_v2.validation import missing_control_space_ids from events.event_handlers import sync_knowledge_fs_bindings_when_app_published_workflow_updated as handler +from models.knowledge_fs import KnowledgeFSControlSpace, KnowledgeFSControlSpaceState +from models.model import App from models.workflow import Workflow +@dataclass +class _RecordingAppBindings: + calls: list[dict[str, object]] = field(default_factory=list) + + def sync_workflow_bindings(self, **kwargs: object) -> list[object]: + self.calls.append(kwargs) + return [] + + +@dataclass +class _Runtime: + app_bindings: _RecordingAppBindings + + +class _FlushRecordingSession(Session): + def __init__(self) -> None: + super().__init__() + self.flush_count = 0 + + def flush(self, objects: Sequence[object] | None = None) -> None: + self.flush_count += 1 + super().flush(objects) + + +def _app() -> App: + return App(id="app-1", tenant_id="tenant-1") + + def _workflow(*space_groups: list[str]) -> Workflow: return Workflow( created_by="account-1", @@ -43,52 +73,53 @@ def test_extracts_deduplicated_control_spaces_in_graph_order() -> None: ) -def test_missing_space_validation_preserves_requested_order() -> None: - session = Mock(spec=Session) - session.scalars.return_value = ["space-b"] +def test_missing_space_validation_preserves_requested_order(sqlite_session: Session) -> None: + active_space = KnowledgeFSControlSpace( + tenant_id="tenant-1", + owner_account_id="account-1", + provisioning_key="provision-b", + knowledge_space_id="knowledge-space-b", + state=KnowledgeFSControlSpaceState.ACTIVE, + ) + active_space.id = "space-b" + sqlite_session.add(active_space) + sqlite_session.flush() workflow = _workflow(["space-a", "space-b", "space-c"]) assert missing_control_space_ids( - session=session, + session=sqlite_session, tenant_id="tenant-1", graph=workflow.graph_dict, ) == ("space-a", "space-c") def test_publish_handler_exactly_syncs_workflow_bindings(monkeypatch: pytest.MonkeyPatch) -> None: - sync = Mock() - monkeypatch.setattr( - handler, - "get_knowledge_fs_runtime", - lambda _session_maker: SimpleNamespace(app_bindings=SimpleNamespace(sync_workflow_bindings=sync)), - ) - app = SimpleNamespace(id="app-1", tenant_id="tenant-1") + bindings = _RecordingAppBindings() + monkeypatch.setattr(handler, "get_knowledge_fs_runtime", lambda _session_maker: _Runtime(bindings)) workflow = _workflow(["space-a", "space-b"]) - publish_session = SimpleNamespace() + publish_session = Session() - handler.handle(app, published_workflow=workflow, session=publish_session) + handler.handle(_app(), published_workflow=workflow, session=publish_session) - sync.assert_called_once_with( - tenant_id="tenant-1", - actor_account_id="account-1", - app_id="app-1", - control_space_ids=["space-a", "space-b"], - session=publish_session, - ) + assert bindings.calls == [ + { + "tenant_id": "tenant-1", + "actor_account_id": "account-1", + "app_id": "app-1", + "control_space_ids": ["space-a", "space-b"], + "session": publish_session, + } + ] def test_publish_union_includes_modern_agent_and_uses_the_publish_transaction(monkeypatch: pytest.MonkeyPatch) -> None: - sync = Mock() - monkeypatch.setattr( - handler, - "get_knowledge_fs_runtime", - lambda _: SimpleNamespace(app_bindings=SimpleNamespace(sync_workflow_bindings=sync)), - ) - session = Mock() + bindings = _RecordingAppBindings() + monkeypatch.setattr(handler, "get_knowledge_fs_runtime", lambda _: _Runtime(bindings)) + publish_session = _FlushRecordingSession() def frozen_spaces(**kwargs: object) -> tuple[str, ...]: - session.flush.assert_called_once() - assert kwargs["session"] is session + assert publish_session.flush_count == 1 + assert kwargs["session"] is publish_session return ("space-b", "space-c") monkeypatch.setattr(handler, "collect_workflow_agent_knowledge_space_ids", frozen_spaces) @@ -96,6 +127,7 @@ def test_publish_union_includes_modern_agent_and_uses_the_publish_transaction(mo graph = workflow.graph_dict graph["nodes"].append({"data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}) workflow.graph = json.dumps(graph) - handler.handle(SimpleNamespace(id="app-1", tenant_id="tenant-1"), published_workflow=workflow, session=session) - assert sync.call_args.kwargs["control_space_ids"] == ["space-a", "space-b", "space-c"] - assert sync.call_args.kwargs["session"] is session + handler.handle(_app(), published_workflow=workflow, session=publish_session) + assert len(bindings.calls) == 1 + assert bindings.calls[0]["control_space_ids"] == ["space-a", "space-b", "space-c"] + assert bindings.calls[0]["session"] is publish_session From 515c9713a9d44d8f94313424855e5fc1de98bccf Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:41:04 +0800 Subject: [PATCH 3/5] test(api): replace test doubles in KnowledgeFS command tests Co-Authored-By: Claude Opus 5.5 --- .../commands/test_knowledge_fs_commands.py | 240 +++++++++++++----- 1 file changed, 178 insertions(+), 62 deletions(-) diff --git a/api/tests/unit_tests/commands/test_knowledge_fs_commands.py b/api/tests/unit_tests/commands/test_knowledge_fs_commands.py index 0c954b4c2d3..1cd107148b2 100644 --- a/api/tests/unit_tests/commands/test_knowledge_fs_commands.py +++ b/api/tests/unit_tests/commands/test_knowledge_fs_commands.py @@ -1,9 +1,9 @@ from __future__ import annotations import json +from dataclasses import dataclass, field from pathlib import Path -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import patch import click import pytest @@ -11,7 +11,15 @@ from click.testing import CliRunner from commands import knowledge_fs as commands from commands.knowledge_fs import _cutover_service, _read_jsonl, knowledge_fs_control_space +from models.knowledge_fs import ( + KnowledgeFSControlSpace, + KnowledgeFSControlSpaceState, + KnowledgeFSLifecycleOperation, + KnowledgeFSLifecycleOutbox, +) from services.knowledge_fs.cleanup import CleanupCompletionReport, CleanupStartReport +from services.knowledge_fs.control_space_commands import KnowledgeFSDeletionIntentResult +from services.knowledge_fs.control_space_management import KnowledgeFSControlSpaceDryRunReport from services.knowledge_fs.cutover import ( CutoverBackfillReport, QuarantineResolutionReport, @@ -20,6 +28,126 @@ from services.knowledge_fs.cutover import ( ) +@dataclass(frozen=True) +class _TenantPayload: + tenant_id: str + + +class _Report: + """Operator report shape: the commands only serialize ``_asdict()``.""" + + def __init__(self, **values: object) -> None: + self._values = values + + def _asdict(self) -> dict[str, object]: + return self._values + + +def _control_space( + control_space_id: str, state: KnowledgeFSControlSpaceState = KnowledgeFSControlSpaceState.ACTIVE +) -> KnowledgeFSControlSpace: + control_space = KnowledgeFSControlSpace( + tenant_id="tenant-1", + owner_account_id="account-1", + provisioning_key="provision-1", + state=state, + ) + control_space.id = control_space_id + return control_space + + +def _deletion_outbox(control_space: KnowledgeFSControlSpace, operation_id: str) -> KnowledgeFSLifecycleOutbox: + return KnowledgeFSLifecycleOutbox( + tenant_id=control_space.tenant_id, + control_space_id=control_space.id, + operation_id=operation_id, + idempotency_key="delete:one", + operation=KnowledgeFSLifecycleOperation.DELETE, + command_payload={}, + expected_control_space_version=0, + ) + + +@dataclass +class _WorkspaceDeletionLifecycleService: + purged_control_spaces: int + preflight_tenants: list[str] = field(default_factory=list) + finalized_tenants: list[str] = field(default_factory=list) + + def assert_workspace_deletion_allowed(self, *, tenant_id: str) -> None: + self.preflight_tenants.append(tenant_id) + + def finalize_workspace_deletion(self, *, tenant_id: str) -> int: + self.finalized_tenants.append(tenant_id) + return self.purged_control_spaces + + +@dataclass +class _ObservationCutoverService: + status_report: dict[str, object] + observation_calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) + status_tenants: list[str] = field(default_factory=list) + + def status(self, *, tenant_id: str) -> dict[str, object]: + self.status_tenants.append(tenant_id) + return self.status_report + + def begin_observation(self, **kwargs: object) -> None: + self.observation_calls.append(("begin", kwargs)) + + def complete_observation(self, **kwargs: object) -> None: + self.observation_calls.append(("complete", kwargs)) + + +@dataclass +class _ManagementService: + registered: KnowledgeFSControlSpace + repaired: KnowledgeFSControlSpace + registrations: list[object] = field(default_factory=list) + repair_calls: list[dict[str, object]] = field(default_factory=list) + dry_run_tenants: list[str | None] = field(default_factory=list) + + def dry_run(self, *, tenant_id: str | None) -> _Report: + self.dry_run_tenants.append(tenant_id) + return _Report(total=1, by_state={"active": 1}) + + def register(self, registration: object) -> tuple[KnowledgeFSControlSpace, bool]: + self.registrations.append(registration) + return self.registered, True + + def repair_registration(self, **kwargs: object) -> KnowledgeFSControlSpace: + self.repair_calls.append(kwargs) + return self.repaired + + +@dataclass +class _CleanupService: + request_calls: list[tuple[object, bool]] = field(default_factory=list) + approve_calls: list[tuple[object, bool]] = field(default_factory=list) + status_calls: list[dict[str, str]] = field(default_factory=list) + + def request_cleanup(self, payload: object, *, apply: bool) -> _Report: + self.request_calls.append((payload, apply)) + return _Report(applied=False) + + def approve_cleanup(self, payload: object, *, apply: bool) -> _Report: + self.approve_calls.append((payload, apply)) + return _Report(applied=True) + + def status(self, *, tenant_id: str, request_id: str) -> dict[str, object]: + self.status_calls.append({"tenant_id": tenant_id, "request_id": request_id}) + return {"state": "approved"} + + +@dataclass +class _OrphanReconciler: + calls: list[dict[str, object]] = field(default_factory=list) + + def reconcile(self, *, limit: int, apply_repairs: bool) -> _Report: + self.calls.append({"limit": limit, "apply_repairs": apply_repairs}) + return _Report(scanned=3, repaired=0) + + def test_control_space_command_exposes_all_p1a_management_operations() -> None: assert { "backfill", @@ -37,18 +165,20 @@ def test_workspace_delete_request_is_dry_run_unless_apply_is_explicit() -> None: patch("commands.knowledge_fs._management_service") as management_factory, patch("commands.knowledge_fs._lifecycle_service") as lifecycle_factory, ): - management_factory.return_value.dry_run.return_value = SimpleNamespace( + management_factory.return_value.dry_run.return_value = KnowledgeFSControlSpaceDryRunReport( total=2, by_state={"active": 2}, + irreversible_deletions=0, ) dry_run = CliRunner().invoke( knowledge_fs_control_space, ["workspace-delete-request", "--tenant-id", "tenant-1"], ) + control_space = _control_space("control-1") lifecycle_factory.return_value.request_workspace_cleanup.return_value = ( - SimpleNamespace( - control_space=SimpleNamespace(id="control-1"), - outbox=SimpleNamespace(operation_id="operation-1"), + KnowledgeFSDeletionIntentResult( + control_space=control_space, + outbox=_deletion_outbox(control_space, "operation-1"), ), ) applied = CliRunner().invoke( @@ -64,9 +194,8 @@ def test_workspace_delete_request_is_dry_run_unless_apply_is_explicit() -> None: def test_workspace_delete_finalize_requires_terminal_preflight_and_explicit_purge_ack() -> None: - with patch("commands.knowledge_fs._lifecycle_service") as lifecycle_factory: - lifecycle_factory.return_value.assert_workspace_deletion_allowed = MagicMock() - lifecycle_factory.return_value.finalize_workspace_deletion.return_value = 2 + lifecycle = _WorkspaceDeletionLifecycleService(purged_control_spaces=2) + with patch("commands.knowledge_fs._lifecycle_service", return_value=lifecycle): preflight = CliRunner().invoke( knowledge_fs_control_space, ["workspace-delete-finalize", "--tenant-id", "tenant-1"], @@ -87,10 +216,10 @@ def test_workspace_delete_finalize_requires_terminal_preflight_and_explicit_purg ) assert preflight.exit_code == 0 - lifecycle_factory.return_value.assert_workspace_deletion_allowed.assert_called_once_with(tenant_id="tenant-1") + assert lifecycle.preflight_tenants == ["tenant-1"] assert denied.exit_code == 2 assert "--acknowledge-control-plane-purge" in denied.output - lifecycle_factory.return_value.finalize_workspace_deletion.assert_called_once_with(tenant_id="tenant-1") + assert lifecycle.finalized_tenants == ["tenant-1"] assert applied.exit_code == 0 assert '"purged_control_spaces": 2' in applied.output @@ -454,35 +583,31 @@ def test_cleanup_complete_cli_requires_explicit_executed_ack_for_apply(tmp_path: assert "--acknowledge-executed" in denied_apply.output -def _report(**values: object) -> SimpleNamespace: - return SimpleNamespace(_asdict=lambda: values) - - def test_operator_commands_delegate_cutover_mutations_with_validated_timestamps( tmp_path: Path, ) -> None: input_path = tmp_path / "operator-input.jsonl" input_path.write_text("{}\n", encoding="utf-8") - payload = SimpleNamespace(tenant_id="tenant-1") - service = MagicMock() - service.status.return_value = { - "legacy_dependency_ready": True, - "open_issues": 0, - "unresolved_cutover_quarantine": 0, - } - for method_name, method in ( - ("inventory", service.inventory), - ("record_shadow_report", service.record_shadow_report), - ("legacy_dependency_dashboard", service.legacy_dependency_dashboard), - ): - method.return_value = _report(operation=method_name) + payload = _TenantPayload(tenant_id="tenant-1") runner = CliRunner() with ( - patch("commands.knowledge_fs._cutover_service", return_value=service), + patch("commands.knowledge_fs._cutover_service") as service_factory, patch("commands.knowledge_fs._read_jsonl", return_value=(payload,)), patch("commands.knowledge_fs._read_one_jsonl", return_value=payload), ): + service = service_factory.return_value + service.status.return_value = { + "legacy_dependency_ready": True, + "open_issues": 0, + "unresolved_cutover_quarantine": 0, + } + for method_name, method in ( + ("inventory", service.inventory), + ("record_shadow_report", service.record_shadow_report), + ("legacy_dependency_dashboard", service.legacy_dependency_dashboard), + ): + method.return_value = _Report(operation=method_name) invocations = [ ["inventory", "--input", str(input_path)], [ @@ -612,8 +737,7 @@ def test_operator_commands_delegate_cutover_mutations_with_validated_timestamps( def test_observation_command_enforces_mutually_exclusive_start_and_completion_modes() -> None: - service = MagicMock() - service.status.return_value = {"phase": "observing"} + service = _ObservationCutoverService(status_report={"phase": "observing"}) runner = CliRunner() base = ["observe", "--tenant-id", "tenant-1", "--expected-cas-version", "7"] @@ -655,12 +779,10 @@ def test_observation_command_enforces_mutually_exclusive_start_and_completion_mo assert "cannot be combined" in conflicting.output assert incomplete.exit_code == 2 assert "observation start requires" in incomplete.output - service.begin_observation.assert_called_once() - service.complete_observation.assert_called_once() + assert [name for name, _ in service.observation_calls] == ["begin", "complete"] def test_legacy_check_returns_a_process_failure_when_any_cutover_gate_is_open() -> None: - service = MagicMock() runner = CliRunner() passing = { "legacy_dependency_ready": True, @@ -673,12 +795,12 @@ def test_legacy_check_returns_a_process_failure_when_any_cutover_gate_is_open() {**passing, "unresolved_cutover_quarantine": 1}, ) + service = _ObservationCutoverService(status_report=passing) with patch("commands.knowledge_fs._cutover_service", return_value=service): - service.status.return_value = passing passed = runner.invoke(knowledge_fs_control_space, ["legacy-check", "--tenant-id", "tenant-1"]) failed = [] for report in failing_reports: - service.status.return_value = report + service.status_report = report failed.append(runner.invoke(knowledge_fs_control_space, ["legacy-check", "--tenant-id", "tenant-1"])) assert passed.exit_code == 0 @@ -692,20 +814,10 @@ def test_management_cleanup_and_reconciliation_commands_preserve_apply_safety( ) -> None: input_path = tmp_path / "operator-input.jsonl" input_path.write_text("{}\n", encoding="utf-8") - payload = SimpleNamespace(tenant_id="tenant-1") - management = MagicMock() - management.dry_run.return_value = _report(total=1, by_state={"active": 1}) - management.register.return_value = (SimpleNamespace(id="control-1"), True) - management.repair_registration.return_value = SimpleNamespace( - id="control-1", - state=SimpleNamespace(value="active"), - ) - cleanup = MagicMock() - cleanup.request_cleanup.return_value = _report(applied=False) - cleanup.approve_cleanup.return_value = _report(applied=True) - cleanup.status.return_value = {"state": "approved"} - reconciler = MagicMock() - reconciler.reconcile.return_value = _report(scanned=3, repaired=0) + payload = _TenantPayload(tenant_id="tenant-1") + management = _ManagementService(registered=_control_space("control-1"), repaired=_control_space("control-1")) + cleanup = _CleanupService() + reconciler = _OrphanReconciler() runner = CliRunner() with ( @@ -752,20 +864,24 @@ def test_management_cleanup_and_reconciliation_commands_preserve_apply_safety( results = [runner.invoke(knowledge_fs_control_space, arguments) for arguments in invocations] assert all(result.exit_code == 0 for result in results), [(result.output, result.exception) for result in results] - registration = management.register.call_args.args[0] + assert management.dry_run_tenants == ["tenant-1"] + assert len(management.registrations) == 1 + registration = management.registrations[0] assert registration.tenant_id == "tenant-1" assert registration.knowledge_space_revision == 2 - management.repair_registration.assert_called_once_with( - tenant_id="tenant-1", - control_space_id="control-1", - expected_resource_version=3, - knowledge_space_id="knowledge-space-1", - knowledge_space_revision=4, - ) - reconciler.reconcile.assert_called_once_with(limit=25, apply_repairs=False) - cleanup.request_cleanup.assert_called_once_with(payload, apply=False) - cleanup.approve_cleanup.assert_called_once_with(payload, apply=True) - cleanup.status.assert_called_once_with(tenant_id="tenant-1", request_id="request-1") + assert management.repair_calls == [ + { + "tenant_id": "tenant-1", + "control_space_id": "control-1", + "expected_resource_version": 3, + "knowledge_space_id": "knowledge-space-1", + "knowledge_space_revision": 4, + } + ] + assert reconciler.calls == [{"limit": 25, "apply_repairs": False}] + assert cleanup.request_calls == [(payload, False)] + assert cleanup.approve_calls == [(payload, True)] + assert cleanup.status_calls == [{"tenant_id": "tenant-1", "request_id": "request-1"}] def test_jsonl_and_timestamp_helpers_reject_ambiguous_operator_input(tmp_path: Path) -> None: From fadf07c43809737a82cd84b61b3d4436b8e3fc9d Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:46:47 +0800 Subject: [PATCH 4/5] test(api): drive KnowledgeFS source task tests through a recording facade Co-Authored-By: Claude Opus 5.5 --- .../tasks/knowledge_fs_task_fakes.py | 185 +++++++ .../test_knowledge_fs_initial_source_tasks.py | 513 ++++++++++-------- .../test_knowledge_fs_source_import_tasks.py | 289 +++++----- 3 files changed, 636 insertions(+), 351 deletions(-) diff --git a/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py b/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py index 86fcca07b64..23eb2d168b8 100644 --- a/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py +++ b/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py @@ -5,12 +5,25 @@ from __future__ import annotations from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from celery.exceptions import Retry + +from core.knowledge_fs.retrieval_contracts import KnowledgeFSPublicFailureResponse from services.knowledge_fs.observability import ( KnowledgeFSBatchStatusMetric, KnowledgeFSCapabilityIssuanceMetric, KnowledgeFSControlSpaceStateMetric, KnowledgeFSLifecycleTaskMetric, ) +from services.knowledge_fs.product_dto import ( + KnowledgeFSSourceConnectionListResponse, + KnowledgeFSSourceConnectionResponse, + KnowledgeFSSourceListResponse, + KnowledgeFSSourceProviderListResponse, + KnowledgeFSSourceProviderResponse, + KnowledgeFSSourceResponse, + KnowledgeFSSourceSyncPolicyResponse, + KnowledgeFSSourceWorkflowResponse, +) @dataclass @@ -40,3 +53,175 @@ class RecordingKnowledgeFSMetrics: if self.gauge_registration_error is not None: raise self.gauge_registration_error self.gauge_readers.append(read_counts) + + +@dataclass +class RecordingRetry: + """Stand-in for a bound Celery ``Task.retry`` that records its keyword arguments.""" + + error: BaseException = field(default_factory=Retry) + calls: list[dict[str, object]] = field(default_factory=list) + + def __call__(self, **kwargs: object) -> Retry: + self.calls.append(kwargs) + raise self.error + + +class RecordingKnowledgeFSFacade: + """Source-facing subset of ``KnowledgeFSDataFacade`` with scripted responses. + + ``script`` sets what a method returns. With several results they are consumed + in order and the last one repeats; an exception result is raised. Calling a + method that was never scripted fails the test instead of returning a dummy. + """ + + def __init__(self) -> None: + self.calls: dict[str, list[dict[str, object]]] = {} + self._results: dict[str, list[object]] = {} + + def script(self, method: str, *results: object) -> None: + self._results[method] = list(results) + + def call_count(self, method: str) -> int: + return len(self.calls.get(method, [])) + + def payload(self, method: str) -> object: + """Return the ``payload`` keyword of the most recent call to ``method``.""" + return self.calls[method][-1]["payload"] + + def _respond(self, method: str, kwargs: dict[str, object]) -> object: + self.calls.setdefault(method, []).append(kwargs) + results = self._results.get(method) + if not results: + raise AssertionError(f"unexpected KnowledgeFS facade call: {method}") + result = results.pop(0) if len(results) > 1 else results[0] + if isinstance(result, BaseException): + raise result + return result + + def list_sources(self, **kwargs: object) -> object: + return self._respond("list_sources", kwargs) + + def get_source(self, **kwargs: object) -> object: + return self._respond("get_source", kwargs) + + def create_source(self, **kwargs: object) -> object: + return self._respond("create_source", kwargs) + + def update_source(self, **kwargs: object) -> object: + return self._respond("update_source", kwargs) + + def list_source_providers(self, **kwargs: object) -> object: + return self._respond("list_source_providers", kwargs) + + def list_source_connections(self, **kwargs: object) -> object: + return self._respond("list_source_connections", kwargs) + + def create_source_connection(self, **kwargs: object) -> object: + return self._respond("create_source_connection", kwargs) + + def import_selected_source_crawl(self, **kwargs: object) -> object: + return self._respond("import_selected_source_crawl", kwargs) + + def import_source_workflow(self, **kwargs: object) -> object: + return self._respond("import_source_workflow", kwargs) + + def consume_namespace_source_preview(self, **kwargs: object) -> object: + return self._respond("consume_namespace_source_preview", kwargs) + + def get_source_workflow(self, **kwargs: object) -> object: + return self._respond("get_source_workflow", kwargs) + + def get_source_sync_policy(self, **kwargs: object) -> object: + return self._respond("get_source_sync_policy", kwargs) + + def update_source_sync_policy(self, **kwargs: object) -> object: + return self._respond("update_source_sync_policy", kwargs) + + +@dataclass(frozen=True) +class FacadeRuntime: + """The ``KnowledgeFSRuntime`` member the source tasks read.""" + + facade: RecordingKnowledgeFSFacade + + +# Response builders construct real DTOs with only the fields a test exercises. +# ``model_construct`` keeps arbitrary remote error codes and messages verbatim. + + +def source_response( + *, + metadata: dict[str, object], + status: str, + version: int, + source_id: str = "source-1", + connection_id: str | None = None, +) -> KnowledgeFSSourceResponse: + return KnowledgeFSSourceResponse.model_construct( + id=source_id, + connection_id=connection_id, + metadata=metadata, + status=status, + version=version, + ) + + +def source_list_response( + *sources: KnowledgeFSSourceResponse, next_cursor: str | None = None +) -> KnowledgeFSSourceListResponse: + return KnowledgeFSSourceListResponse.model_construct(data=list(sources), next_cursor=next_cursor) + + +def source_workflow_response( + *, + workflow_id: str, + state: str, + source_id: str | None = None, + failure_message: str | None = None, + last_error_code: str | None = None, +) -> KnowledgeFSSourceWorkflowResponse: + failure = ( + None if failure_message is None else KnowledgeFSPublicFailureResponse.model_construct(message=failure_message) + ) + return KnowledgeFSSourceWorkflowResponse.model_construct( + id=workflow_id, + state=state, + source_id=source_id, + failure=failure, + last_error_code=last_error_code, + ) + + +def sync_policy_response(*, revision: int) -> KnowledgeFSSourceSyncPolicyResponse: + return KnowledgeFSSourceSyncPolicyResponse.model_construct(revision=revision) + + +def source_provider_list_response(*available_provider_ids: str) -> KnowledgeFSSourceProviderListResponse: + return KnowledgeFSSourceProviderListResponse.model_construct( + data=[ + KnowledgeFSSourceProviderResponse.model_construct(id=provider_id, available=True) + for provider_id in available_provider_ids + ] + ) + + +def source_connection_response( + *, + connection_id: str, + provider_id: str, + status: str, + configuration: dict[str, bool | int | str] | None = None, +) -> KnowledgeFSSourceConnectionResponse: + return KnowledgeFSSourceConnectionResponse.model_construct( + id=connection_id, + provider_id=provider_id, + status=status, + configuration=configuration or {}, + ) + + +def source_connection_list_response( + *connections: KnowledgeFSSourceConnectionResponse, next_cursor: str | None = None +) -> KnowledgeFSSourceConnectionListResponse: + return KnowledgeFSSourceConnectionListResponse.model_construct(data=list(connections), next_cursor=next_cursor) diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py index b967769d76c..036f5f84766 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py @@ -1,11 +1,14 @@ from collections.abc import Generator from contextlib import contextmanager -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from dataclasses import dataclass, field +from unittest.mock import patch import pytest -from models.knowledge_fs import KnowledgeFSControlSpaceState +from core.db.session_factory import session_factory +from models.account import Account +from models.knowledge_fs import KnowledgeFSControlSpace, KnowledgeFSControlSpaceState +from models.oauth import DatasourceProvider from services.knowledge_fs.product_dto import ( KnowledgeFSInitialOnlineDocumentSourcePayload, KnowledgeFSInitialOnlineDriveSourcePayload, @@ -25,8 +28,17 @@ from tasks.knowledge_fs_initial_source_tasks import ( start_initial_website_source_import, submit_initial_source_for_upgrade, ) - -_DEFAULT_CREDENTIAL = object() +from tests.unit_tests.tasks.knowledge_fs_task_fakes import ( + FacadeRuntime, + RecordingKnowledgeFSFacade, + source_connection_list_response, + source_connection_response, + source_list_response, + source_provider_list_response, + source_response, + source_workflow_response, + sync_policy_response, +) def _payload(sync_policy: str = "daily") -> KnowledgeFSInitialWebsiteSourcePayload: @@ -111,102 +123,149 @@ def _drive_payload() -> KnowledgeFSInitialOnlineDriveSourcePayload: ) -def _facade() -> MagicMock: - facade = MagicMock() - facade.list_sources.return_value = SimpleNamespace(data=[], next_cursor=None) - facade.list_source_providers.return_value = SimpleNamespace( - data=[SimpleNamespace(id="plugin-daemon-website", available=True)] - ) - facade.list_source_connections.return_value = SimpleNamespace( - data=[ - SimpleNamespace( - configuration={ - "credentialId": "firecrawl-credential-1", - "datasource": "crawl", - "pluginId": "langgenius/firecrawl_datasource", - "provider": "firecrawl", - "providerKind": "website", - }, - id="connection-1", +_FIRECRAWL_CONNECTION_CONFIGURATION: dict[str, bool | int | str] = { + "credentialId": "firecrawl-credential-1", + "datasource": "crawl", + "pluginId": "langgenius/firecrawl_datasource", + "provider": "firecrawl", + "providerKind": "website", +} + + +@dataclass(frozen=True) +class _Credential: + id: str + name: str + provider: str = "firecrawl" + plugin_id: str = "langgenius/firecrawl_datasource" + + +_FIRECRAWL_CREDENTIAL = _Credential(id="firecrawl-credential-1", name="Firecrawl") + + +@dataclass +class _ControlSpaceRepository: + control_space: KnowledgeFSControlSpace | None + lookups: list[dict[str, str]] = field(default_factory=list) + + def get(self, *, tenant_id: str, control_space_id: str) -> KnowledgeFSControlSpace | None: + self.lookups.append({"tenant_id": tenant_id, "control_space_id": control_space_id}) + return self.control_space + + +def _facade() -> RecordingKnowledgeFSFacade: + facade = RecordingKnowledgeFSFacade() + facade.script("list_sources", source_list_response()) + facade.script("list_source_providers", source_provider_list_response("plugin-daemon-website")) + facade.script( + "list_source_connections", + source_connection_list_response( + source_connection_response( + configuration=_FIRECRAWL_CONNECTION_CONFIGURATION, + connection_id="connection-1", provider_id="plugin-daemon-website", status="active", ) - ], - next_cursor=None, + ), ) - facade.create_source_connection.return_value = SimpleNamespace( - configuration={"credentialId": "firecrawl-credential-1"}, - id="connection-created", - provider_id="plugin-daemon-website", - status="active", + facade.script( + "create_source_connection", + source_connection_response( + configuration={"credentialId": "firecrawl-credential-1"}, + connection_id="connection-created", + provider_id="plugin-daemon-website", + status="active", + ), ) - facade.create_source.return_value = SimpleNamespace(id="source-1") - facade.import_selected_source_crawl.return_value = SimpleNamespace( - id="workflow-1", - source_id="source-1", - state="completed", + facade.script("create_source", source_response(metadata={}, status="disabled", version=1)) + facade.script( + "import_selected_source_crawl", + source_workflow_response(workflow_id="workflow-1", source_id="source-1", state="completed"), ) - facade.get_source_workflow.return_value = SimpleNamespace( - id="workflow-1", - source_id="source-1", - state="completed", + facade.script( + "get_source_workflow", + source_workflow_response(workflow_id="workflow-1", source_id="source-1", state="completed"), ) - facade.get_source.return_value = SimpleNamespace( - metadata={ - "clientRequestId": "initial-website-source:operation-1", - "parameters": {"limit": 99}, - "preview": True, - }, - status="disabled", - version=3, + facade.script( + "get_source", + source_response( + metadata={ + "clientRequestId": "initial-website-source:operation-1", + "parameters": {"limit": 99}, + "preview": True, + }, + status="disabled", + version=3, + ), ) - facade.update_source.return_value = SimpleNamespace( - metadata={ - "clientRequestId": "initial-website-source:operation-1", - "preview": False, - }, - status="active", - version=4, + facade.script( + "update_source", + source_response( + metadata={ + "clientRequestId": "initial-website-source:operation-1", + "preview": False, + }, + status="active", + version=4, + ), ) - facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("not found") + facade.script("get_source_sync_policy", KnowledgeFSProductResourceNotFoundError("not found")) + facade.script("update_source_sync_policy", None) return facade +def _seed_account_and_credential(credential: _Credential | None) -> None: + account = Account(name="Initial source owner", email="initial-source-owner@example.com") + account.id = "account-1" + with session_factory.create_session() as session: + session.add(account) + if credential is not None: + provider = DatasourceProvider( + tenant_id="tenant-1", + name=credential.name, + provider=credential.provider, + plugin_id=credential.plugin_id, + auth_type="api_key", + encrypted_credentials={}, + ) + provider.id = credential.id + session.add(provider) + session.commit() + + @contextmanager def _runtime( - facade: MagicMock, + facade: RecordingKnowledgeFSFacade, *, - credential: SimpleNamespace | None | object = _DEFAULT_CREDENTIAL, + credential: _Credential | None = _FIRECRAWL_CREDENTIAL, state: KnowledgeFSControlSpaceState = KnowledgeFSControlSpaceState.ACTIVE, knowledge_space_id: str | None = "space-1", -) -> Generator[MagicMock]: - session_context = MagicMock() - session = MagicMock() - session.scalar.return_value = ( - SimpleNamespace(id="firecrawl-credential-1", name="Firecrawl") - if credential is _DEFAULT_CREDENTIAL - else credential + control_space_found: bool = True, +) -> Generator[_ControlSpaceRepository]: + _seed_account_and_credential(credential) + control_space = KnowledgeFSControlSpace( + tenant_id="tenant-1", + owner_account_id="account-1", + provisioning_key="provision-1", + state=state, + knowledge_space_id=knowledge_space_id, ) - session_context.__enter__.return_value = session - session_maker = MagicMock(return_value=session_context) + repository = _ControlSpaceRepository(control_space if control_space_found else None) with ( patch( - "tasks.knowledge_fs_initial_source_tasks.session_factory.get_session_maker", - return_value=session_maker, + "tasks.knowledge_fs_initial_source_tasks.SQLAlchemyKnowledgeFSControlSpaceRepository", + return_value=repository, + ), + patch( + "tasks.knowledge_fs_initial_source_tasks.get_knowledge_fs_runtime", + return_value=FacadeRuntime(facade), ), - patch("tasks.knowledge_fs_initial_source_tasks.SQLAlchemyKnowledgeFSControlSpaceRepository") as repository_type, - patch("tasks.knowledge_fs_initial_source_tasks.get_knowledge_fs_runtime") as get_runtime, ): - repository_type.return_value.get.return_value = SimpleNamespace( - state=state, - knowledge_space_id=knowledge_space_id, - ) - get_runtime.return_value.facade = facade - yield repository_type + yield repository def _start( - facade: MagicMock, + facade: RecordingKnowledgeFSFacade, payload: KnowledgeFSInitialWebsiteSourcePayload, *, workflow_id: str | None = None, @@ -227,7 +286,7 @@ def test_initial_website_source_import_recrawls_exact_selection_and_configures_d assert _start(facade, _payload()) == "workflow-1" - source_payload = facade.create_source.call_args.kwargs["payload"] + source_payload = facade.payload("create_source") assert source_payload.status == "disabled" assert source_payload.connection_id == "connection-1" assert source_payload.metadata["datasourceParameterMode"] == "exact" @@ -240,16 +299,16 @@ def test_initial_website_source_import_recrawls_exact_selection_and_configures_d "limit": 25, "url": "https://docs.dify.ai", } - import_payload = facade.import_selected_source_crawl.call_args.kwargs["payload"] + import_payload = facade.payload("import_selected_source_crawl") assert import_payload.source_urls == [ "https://docs.dify.ai/a", "https://docs.dify.ai/b", ] - source_update_payload = facade.update_source.call_args.kwargs["payload"] + source_update_payload = facade.payload("update_source") assert source_update_payload.expected_version == 3 assert source_update_payload.metadata == {"preview": False} assert source_update_payload.status == "active" - sync_payload = facade.update_source_sync_policy.call_args.kwargs["payload"] + sync_payload = facade.payload("update_source_sync_policy") assert sync_payload.mode == "interval" assert sync_payload.enabled is True assert sync_payload.expected_revision == 0 @@ -275,7 +334,7 @@ def test_initial_website_source_import_preserves_distinct_urls_with_the_same_can assert _start(facade, payload) == "workflow-1" - import_payload = facade.import_selected_source_crawl.call_args.kwargs["payload"] + import_payload = facade.payload("import_selected_source_crawl") assert import_payload.source_urls == [ "https://docs.dify.ai/a/", "https://DOCS.dify.ai:443/a#section", @@ -287,7 +346,7 @@ def test_initial_website_source_import_preserves_legacy_crawl_option_projection( assert _start(facade, _legacy_payload()) == "workflow-1" - source_payload = facade.create_source.call_args.kwargs["payload"] + source_payload = facade.payload("create_source") assert "datasourceParameterMode" not in source_payload.metadata assert "parameters" not in source_payload.metadata assert source_payload.metadata["crawlOptions"] == { @@ -334,25 +393,38 @@ def test_initial_connector_source_import_uses_exact_binding_and_selection( expected_provider_name: str, ) -> None: facade = _facade() - facade.list_source_providers.return_value = SimpleNamespace(data=[SimpleNamespace(id=provider_id, available=True)]) - facade.list_source_connections.return_value = SimpleNamespace(data=[], next_cursor=None) - facade.create_source_connection.return_value = SimpleNamespace(id="connector-1", status="active") - facade.import_source_workflow.return_value = SimpleNamespace( - id="connector-workflow-1", - source_id="source-1", - state="completed", + facade.script("list_source_providers", source_provider_list_response(provider_id)) + facade.script("list_source_connections", source_connection_list_response()) + facade.script( + "create_source_connection", + source_connection_response(connection_id="connector-1", provider_id=provider_id, status="active"), ) - facade.get_source.return_value = SimpleNamespace( - metadata={"clientRequestId": "initial-source:operation-1", "preview": True}, - status="disabled", - version=3, + facade.script( + "import_source_workflow", + source_workflow_response(workflow_id="connector-workflow-1", source_id="source-1", state="completed"), ) - facade.update_source.return_value = SimpleNamespace( - metadata={"clientRequestId": "initial-source:operation-1", "preview": False}, - status="active", - version=4, + facade.script( + "get_source", + source_response( + metadata={"clientRequestId": "initial-source:operation-1", "preview": True}, + status="disabled", + version=3, + ), + ) + facade.script( + "update_source", + source_response( + metadata={"clientRequestId": "initial-source:operation-1", "preview": False}, + status="active", + version=4, + ), + ) + credential = _Credential( + id=credential_id, + name=expected_source_name, + provider=payload.provider, + plugin_id=payload.plugin_id, ) - credential = SimpleNamespace(id=credential_id, name=expected_source_name) with _runtime(facade, credential=credential): result = start_initial_source_import( @@ -364,7 +436,7 @@ def test_initial_connector_source_import_uses_exact_binding_and_selection( ) assert result == "connector-workflow-1" - connection_payload = facade.create_source_connection.call_args.kwargs["payload"] + connection_payload = facade.payload("create_source_connection") assert connection_payload.provider_id == provider_id assert connection_payload.configuration == { "credentialId": credential_id, @@ -373,45 +445,50 @@ def test_initial_connector_source_import_uses_exact_binding_and_selection( "provider": payload.provider, "providerKind": "online-document" if workflow_kind == "online-document-import" else "online-drive", } - source_payload = facade.create_source.call_args.kwargs["payload"] + source_payload = facade.payload("create_source") assert source_payload.connection_id == "connector-1" assert source_payload.name == expected_source_name assert source_payload.type == "connector" assert source_payload.uri == "connector://connector-1" assert source_payload.metadata["providerId"] == provider_id assert source_payload.metadata["providerName"] == expected_provider_name - import_payload = facade.import_source_workflow.call_args.kwargs["payload"].root + import_payload = facade.payload("import_source_workflow").root assert import_payload.kind == workflow_kind assert import_payload.items == payload.selection - sync_payload = facade.update_source_sync_policy.call_args.kwargs["payload"] + sync_payload = facade.payload("update_source_sync_policy") assert sync_payload.mode == ("manual" if workflow_kind == "online-drive-import" else "interval") def test_initial_website_source_import_reuses_source_across_pages_and_preserves_failure() -> None: facade = _facade() - existing_source = SimpleNamespace( - id="existing-source", + existing_source = source_response( + source_id="existing-source", metadata={"clientRequestId": "initial-website-source:operation-1"}, + status="disabled", version=2, ) - facade.list_sources.side_effect = [ - SimpleNamespace( - data=[SimpleNamespace(id="other", metadata={})], + facade.script( + "list_sources", + source_list_response( + source_response(source_id="other", metadata={}, status="active", version=1), next_cursor="next-source", ), - SimpleNamespace(data=[existing_source], next_cursor=None), - ] - facade.import_selected_source_crawl.return_value = SimpleNamespace( - failure=SimpleNamespace(message="Source document materialization failed"), - id="failed-workflow", - last_error_code="SOURCE_DOCUMENT_MATERIALIZATION_FAILED", - state="failed", + source_list_response(existing_source), + ) + facade.script( + "import_selected_source_crawl", + source_workflow_response( + workflow_id="failed-workflow", + state="failed", + failure_message="Source document materialization failed", + last_error_code="SOURCE_DOCUMENT_MATERIALIZATION_FAILED", + ), ) assert _start(facade, _payload()) == "failed-workflow" - facade.create_source.assert_not_called() - facade.list_source_connections.assert_not_called() - source_update_payload = facade.update_source.call_args.kwargs["payload"] + assert facade.call_count("create_source") == 0 + assert facade.call_count("list_source_connections") == 0 + source_update_payload = facade.payload("update_source") assert source_update_payload.status == "disabled" assert "parameters" not in source_update_payload.metadata assert source_update_payload.metadata["preview"] is False @@ -421,12 +498,12 @@ def test_initial_website_source_import_reuses_source_across_pages_and_preserves_ "enabled": True, "mode": "interval", } - facade.update_source_sync_policy.assert_not_called() + assert facade.call_count("update_source_sync_policy") == 0 def test_initial_website_source_import_exposes_preview_consume_rejection() -> None: facade = _facade() - facade.consume_namespace_source_preview.side_effect = KnowledgeFSProductRequestRejectedError(status_code=409) + facade.script("consume_namespace_source_preview", KnowledgeFSProductRequestRejectedError(status_code=409)) payload = _payload().model_copy( update={ "credential_id": "firecrawl-credential-1", @@ -445,7 +522,7 @@ def test_initial_website_source_import_exposes_preview_consume_rejection() -> No with pytest.raises(KnowledgeFSProductRequestRejectedError): _start(facade, payload) - source_update = facade.update_source.call_args.kwargs["payload"] + source_update = facade.payload("update_source") assert source_update.status == "disabled" assert source_update.metadata["preview"] is False assert source_update.metadata["initialImport"]["state"] == "failed" @@ -472,34 +549,25 @@ def test_initial_website_source_import_configures_remaining_sync_modes( expected_custom_interval_seconds: int | None, ) -> None: facade = _facade() - facade.list_source_connections.side_effect = [ - SimpleNamespace( - data=[SimpleNamespace(id="inactive", provider_id="other", status="disabled")], + facade.script( + "list_source_connections", + source_connection_list_response( + source_connection_response(connection_id="inactive", provider_id="other", status="revoked"), next_cursor="next-connection", ), - SimpleNamespace( - data=[ - SimpleNamespace( - configuration={ - "credentialId": "firecrawl-credential-1", - "datasource": "crawl", - "pluginId": "langgenius/firecrawl_datasource", - "provider": "firecrawl", - "providerKind": "website", - }, - id="connection-2", - provider_id="plugin-daemon-website", - status="active", - ) - ], - next_cursor=None, + source_connection_list_response( + source_connection_response( + configuration=_FIRECRAWL_CONNECTION_CONFIGURATION, + connection_id="connection-2", + provider_id="plugin-daemon-website", + status="active", + ) ), - ] - facade.get_source_sync_policy.side_effect = None - facade.get_source_sync_policy.return_value = SimpleNamespace(revision=7) + ) + facade.script("get_source_sync_policy", sync_policy_response(revision=7)) assert _start(facade, _payload(sync_policy)) == "workflow-1" - sync_payload = facade.update_source_sync_policy.call_args.kwargs["payload"] + sync_payload = facade.payload("update_source_sync_policy") assert sync_payload.mode == expected_mode assert sync_payload.enabled is expected_enabled assert sync_payload.custom_interval_seconds == expected_custom_interval_seconds @@ -509,19 +577,22 @@ def test_initial_website_source_import_configures_remaining_sync_modes( def test_initial_website_source_import_does_not_recommit_active_source_after_retry() -> None: facade = _facade() - facade.get_source.return_value = SimpleNamespace( - metadata={ - "clientRequestId": "initial-website-source:operation-1", - "preview": False, - }, - status="active", - version=4, + facade.script( + "get_source", + source_response( + metadata={ + "clientRequestId": "initial-website-source:operation-1", + "preview": False, + }, + status="active", + version=4, + ), ) assert _start(facade, _payload()) == "workflow-1" - facade.update_source.assert_not_called() - sync_payload = facade.update_source_sync_policy.call_args.kwargs["payload"] + assert facade.call_count("update_source") == 0 + sync_payload = facade.payload("update_source_sync_policy") assert sync_payload.expected_source_version == 4 @@ -570,8 +641,7 @@ def test_initial_website_source_import_rejects_unavailable_spaces( def test_initial_website_source_import_rejects_missing_control_space() -> None: facade = _facade() - with _runtime(facade) as repository_type: - repository_type.return_value.get.return_value = None + with _runtime(facade, control_space_found=False): with pytest.raises(RuntimeError, match="control-space was not found"): start_initial_website_source_import( tenant_id="tenant-1", @@ -584,7 +654,7 @@ def test_initial_website_source_import_rejects_missing_control_space() -> None: def test_initial_website_source_import_rejects_unavailable_provider() -> None: unavailable_provider = _facade() - unavailable_provider.list_source_providers.return_value = SimpleNamespace(data=[]) + unavailable_provider.script("list_source_providers", source_provider_list_response()) with pytest.raises(RuntimeError, match="provider is unavailable"): _start(unavailable_provider, _payload()) @@ -603,20 +673,22 @@ def test_initial_website_source_import_rejects_missing_firecrawl_credential() -> payload=_payload(), ) - facade.create_source_connection.assert_not_called() - facade.create_source.assert_not_called() + assert facade.call_count("create_source_connection") == 0 + assert facade.call_count("create_source") == 0 def test_initial_website_source_import_creates_missing_firecrawl_connection() -> None: facade = _facade() - facade.list_source_connections.return_value = SimpleNamespace( - data=[SimpleNamespace(id="inactive", provider_id="other", status="disabled")], - next_cursor=None, + facade.script( + "list_source_connections", + source_connection_list_response( + source_connection_response(connection_id="inactive", provider_id="other", status="revoked") + ), ) assert _start(facade, _payload()) == "workflow-1" - connection_payload = facade.create_source_connection.call_args.kwargs["payload"] + connection_payload = facade.payload("create_source_connection") assert connection_payload.auth_kind == "endpoint" assert connection_payload.configuration == { "credentialId": "firecrawl-credential-1", @@ -627,43 +699,36 @@ def test_initial_website_source_import_creates_missing_firecrawl_connection() -> } assert connection_payload.credentials == {} assert connection_payload.provider_id == "plugin-daemon-website" - source_payload = facade.create_source.call_args.kwargs["payload"] + source_payload = facade.payload("create_source") assert source_payload.connection_id == "connection-created" def test_initial_website_source_import_retries_provisioning_firecrawl_connection() -> None: facade = _facade() - facade.list_source_connections.return_value = SimpleNamespace( - data=[ - SimpleNamespace( - configuration={ - "credentialId": "firecrawl-credential-1", - "datasource": "crawl", - "pluginId": "langgenius/firecrawl_datasource", - "provider": "firecrawl", - "providerKind": "website", - }, - id="connection-1", + facade.script( + "list_source_connections", + source_connection_list_response( + source_connection_response( + configuration=_FIRECRAWL_CONNECTION_CONFIGURATION, + connection_id="connection-1", provider_id="plugin-daemon-website", status="provisioning", ) - ], - next_cursor=None, + ), ) with pytest.raises(KnowledgeFSInitialSourceNotReadyError, match="still provisioning"): _start(facade, _payload()) - facade.create_source_connection.assert_not_called() - facade.create_source.assert_not_called() + assert facade.call_count("create_source_connection") == 0 + assert facade.call_count("create_source") == 0 def test_initial_website_source_import_retries_running_workflow() -> None: facade = _facade() - facade.import_selected_source_crawl.return_value = SimpleNamespace( - id="running-workflow", - source_id="source-1", - state="running", + facade.script( + "import_selected_source_crawl", + source_workflow_response(workflow_id="running-workflow", source_id="source-1", state="running"), ) with pytest.raises(KnowledgeFSInitialSourceNotReadyError, match="still running") as raised: @@ -677,49 +742,53 @@ def test_initial_website_source_import_polls_existing_workflow_without_recreatin assert _start(facade, _payload(), workflow_id="workflow-1") == "workflow-1" - facade.get_source_workflow.assert_called_once_with( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - run_id="workflow-1", - ) - facade.list_sources.assert_not_called() - facade.list_source_connections.assert_not_called() - facade.create_source.assert_not_called() - facade.import_selected_source_crawl.assert_not_called() - source_update_payload = facade.update_source.call_args.kwargs["payload"] + assert facade.calls["get_source_workflow"] == [ + { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "control-1", + "run_id": "workflow-1", + } + ] + assert facade.call_count("list_sources") == 0 + assert facade.call_count("list_source_connections") == 0 + assert facade.call_count("create_source") == 0 + assert facade.call_count("import_selected_source_crawl") == 0 + source_update_payload = facade.payload("update_source") assert source_update_payload.status == "active" def test_initial_website_source_import_keeps_polling_existing_running_workflow() -> None: facade = _facade() - facade.get_source_workflow.return_value = SimpleNamespace( - id="workflow-1", - source_id="source-1", - state="importing", + facade.script( + "get_source_workflow", + source_workflow_response(workflow_id="workflow-1", source_id="source-1", state="importing"), ) with pytest.raises(KnowledgeFSInitialSourceNotReadyError, match="still running") as raised: _start(facade, _payload(), workflow_id="workflow-1") assert raised.value.workflow_id == "workflow-1" - facade.import_selected_source_crawl.assert_not_called() - facade.get_source.assert_not_called() + assert facade.call_count("import_selected_source_crawl") == 0 + assert facade.call_count("get_source") == 0 def test_initial_website_source_import_exposes_failed_source_without_activating_it() -> None: facade = _facade() - facade.get_source_workflow.return_value = SimpleNamespace( - failure=SimpleNamespace(message="Source document materialization failed"), - id="workflow-1", - last_error_code="SOURCE_DOCUMENT_MATERIALIZATION_FAILED", - source_id="source-1", - state="failed", + facade.script( + "get_source_workflow", + source_workflow_response( + workflow_id="workflow-1", + source_id="source-1", + state="failed", + failure_message="Source document materialization failed", + last_error_code="SOURCE_DOCUMENT_MATERIALIZATION_FAILED", + ), ) assert _start(facade, _payload(), workflow_id="workflow-1") == "workflow-1" - source_update_payload = facade.update_source.call_args.kwargs["payload"] + source_update_payload = facade.payload("update_source") assert source_update_payload.status == "disabled" assert source_update_payload.metadata["preview"] is False initial_import = source_update_payload.metadata["initialImport"] @@ -735,7 +804,7 @@ def test_initial_website_source_import_exposes_failed_source_without_activating_ "mode": "interval", } assert len(initial_import["configurationFingerprint"]) == 64 - facade.update_source_sync_policy.assert_not_called() + assert facade.call_count("update_source_sync_policy") == 0 def test_initial_website_source_task_returns_result_and_retries_not_ready_error() -> None: @@ -904,21 +973,25 @@ def test_initial_source_task_does_not_retry_authoritative_missing_resource() -> def test_upgrade_source_submission_succeeds_when_first_import_submission_fails() -> None: facade = _facade() - facade.create_source.return_value = SimpleNamespace( - connection_id="connection-1", - id="source-1", - metadata={"clientRequestId": "initial-website-source:operation-1", "preview": True}, - status="disabled", - version=1, + facade.script( + "create_source", + source_response( + connection_id="connection-1", + metadata={"clientRequestId": "initial-website-source:operation-1", "preview": True}, + status="disabled", + version=1, + ), ) - facade.update_source.return_value = SimpleNamespace( - connection_id="connection-1", - id="source-1", - metadata={"clientRequestId": "initial-website-source:operation-1", "preview": False}, - status="active", - version=2, + facade.script( + "update_source", + source_response( + connection_id="connection-1", + metadata={"clientRequestId": "initial-website-source:operation-1", "preview": False}, + status="active", + version=2, + ), ) - facade.import_selected_source_crawl.side_effect = RuntimeError("crawl queue unavailable") + facade.script("import_selected_source_crawl", RuntimeError("crawl queue unavailable")) with _runtime(facade): result = submit_initial_source_for_upgrade( @@ -933,4 +1006,4 @@ def test_upgrade_source_submission_succeeds_when_first_import_submission_fails() assert result.source_id == "source-1" assert result.workflow_id is None assert result.workflow_error == "RuntimeError" - facade.update_source_sync_policy.assert_called_once() + assert facade.call_count("update_source_sync_policy") == 1 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 88bd865f6bf..561a7db5739 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 @@ -1,10 +1,10 @@ -from types import SimpleNamespace -from typing import Literal -from unittest.mock import MagicMock, patch +from typing import Literal, Never +from unittest.mock import patch import pytest from celery.exceptions import Retry +from services.knowledge_fs.product_dto import KnowledgeFSSourceResponse, KnowledgeFSSourceWorkflowResponse from services.knowledge_fs.product_remote import ( KnowledgeFSProductRequestRejectedError, KnowledgeFSProductResourceNotFoundError, @@ -14,11 +14,17 @@ from tasks.knowledge_fs_source_import_tasks import ( KnowledgeFSSourceImportNotReadyError, finalize_source_import_once, ) +from tests.unit_tests.tasks.knowledge_fs_task_fakes import ( + FacadeRuntime, + RecordingKnowledgeFSFacade, + RecordingRetry, + source_response, + source_workflow_response, +) -def _source(*, status: str = "syncing") -> SimpleNamespace: - return SimpleNamespace( - id="source-1", +def _source(*, status: str = "syncing") -> KnowledgeFSSourceResponse: + return source_response( metadata={ "parameters": {"limit": 99}, "preview": False, @@ -34,9 +40,23 @@ def _source(*, status: str = "syncing") -> SimpleNamespace: ) -def _run(facade: MagicMock) -> str: - with patch("tasks.knowledge_fs_source_import_tasks.get_knowledge_fs_runtime") as runtime: - runtime.return_value.facade = facade +def _completed_workflow() -> KnowledgeFSSourceWorkflowResponse: + return source_workflow_response(workflow_id="import-1", state="completed") + + +def _facade(workflow: KnowledgeFSSourceWorkflowResponse) -> RecordingKnowledgeFSFacade: + facade = RecordingKnowledgeFSFacade() + facade.script("get_source_workflow", workflow) + facade.script("get_source_sync_policy", KnowledgeFSProductResourceNotFoundError("missing")) + facade.script("update_source_sync_policy", None) + return facade + + +def _run(facade: RecordingKnowledgeFSFacade) -> str: + with patch( + "tasks.knowledge_fs_source_import_tasks.get_knowledge_fs_runtime", + return_value=FacadeRuntime(facade), + ): return finalize_source_import_once( tenant_id="tenant-1", account_id="account-1", @@ -46,45 +66,48 @@ def _run(facade: MagicMock) -> str: ) +def _raise(error: Exception) -> Never: + raise error + + def test_finalize_source_import_waits_for_terminal_workflow() -> None: - facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace(state="importing") + facade = _facade(source_workflow_response(workflow_id="import-1", state="importing")) with pytest.raises(KnowledgeFSSourceImportNotReadyError): _run(facade) - facade.get_source.assert_not_called() + assert facade.call_count("get_source") == 0 def test_finalize_source_import_activates_source_then_applies_policy_for_new_version() -> None: - facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace(id="import-1", state="completed") - facade.get_source.return_value = _source() - facade.update_source.return_value = SimpleNamespace( - id="source-1", - metadata={ - "lastImport": { - "kind": "crawl-preview-selection", - "previewWorkflowId": "preview-1", - "state": "completed", - "syncPolicy": {"enabled": False, "mode": "manual"}, - "workflowId": "import-1", + facade = _facade(_completed_workflow()) + facade.script("get_source", _source()) + facade.script( + "update_source", + source_response( + metadata={ + "lastImport": { + "kind": "crawl-preview-selection", + "previewWorkflowId": "preview-1", + "state": "completed", + "syncPolicy": {"enabled": False, "mode": "manual"}, + "workflowId": "import-1", + }, + "preview": False, }, - "preview": False, - }, - status="active", - version=5, + status="active", + version=5, + ), ) - facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("missing") assert _run(facade) == "import-1" - policy = facade.update_source_sync_policy.call_args.kwargs["payload"] + policy = facade.payload("update_source_sync_policy") assert policy.enabled is False assert policy.mode == "manual" assert policy.expected_revision == 0 assert policy.expected_source_version == 5 - update = facade.update_source.call_args.kwargs["payload"] + update = facade.payload("update_source") assert update.status == "active" assert "parameters" not in update.metadata assert update.metadata["pendingImport"] is None @@ -92,94 +115,98 @@ def test_finalize_source_import_activates_source_then_applies_policy_for_new_ver def test_finalize_source_import_retries_policy_after_source_activation() -> None: - facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace(id="import-1", state="completed") - facade.get_source.return_value = SimpleNamespace( - id="source-1", - metadata={ - "lastImport": { - "kind": "crawl-preview-selection", - "previewWorkflowId": "preview-1", - "state": "completed", - "syncPolicy": {"enabled": True, "mode": "interval"}, - "workflowId": "import-1", + facade = _facade(_completed_workflow()) + facade.script( + "get_source", + source_response( + metadata={ + "lastImport": { + "kind": "crawl-preview-selection", + "previewWorkflowId": "preview-1", + "state": "completed", + "syncPolicy": {"enabled": True, "mode": "interval"}, + "workflowId": "import-1", + }, + "preview": False, }, - "preview": False, - }, - status="active", - version=5, + status="active", + version=5, + ), ) - facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("missing") assert _run(facade) == "import-1" - facade.update_source.assert_not_called() - policy = facade.update_source_sync_policy.call_args.kwargs["payload"] + assert facade.call_count("update_source") == 0 + policy = facade.payload("update_source_sync_policy") assert policy.expected_source_version == 5 def test_finalize_source_import_recovers_a_stale_failed_marker_after_workflow_completion() -> None: - facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace(id="import-1", state="completed") - facade.get_source.return_value = SimpleNamespace( - id="source-1", - metadata={ - "lastImport": { - "errorCode": "SOURCE_OPERATION_FAILED", - "errorMessage": "provider failed", - "kind": "crawl-preview-selection", - "previewWorkflowId": "preview-1", - "state": "failed", - "syncPolicy": {"enabled": False, "mode": "manual"}, - "workflowId": "import-1", + facade = _facade(_completed_workflow()) + facade.script( + "get_source", + source_response( + metadata={ + "lastImport": { + "errorCode": "SOURCE_OPERATION_FAILED", + "errorMessage": "provider failed", + "kind": "crawl-preview-selection", + "previewWorkflowId": "preview-1", + "state": "failed", + "syncPolicy": {"enabled": False, "mode": "manual"}, + "workflowId": "import-1", + }, + "pendingImport": None, + "preview": False, }, - "pendingImport": None, - "preview": False, - }, - status="error", - version=5, + status="error", + version=5, + ), ) - facade.update_source.return_value = SimpleNamespace( - id="source-1", - metadata={ - "lastImport": { - "kind": "crawl-preview-selection", - "previewWorkflowId": "preview-1", - "state": "completed", - "syncPolicy": {"enabled": False, "mode": "manual"}, - "workflowId": "import-1", + facade.script( + "update_source", + source_response( + metadata={ + "lastImport": { + "kind": "crawl-preview-selection", + "previewWorkflowId": "preview-1", + "state": "completed", + "syncPolicy": {"enabled": False, "mode": "manual"}, + "workflowId": "import-1", + }, + "pendingImport": None, + "preview": False, }, - "pendingImport": None, - "preview": False, - }, - status="active", - version=6, + status="active", + version=6, + ), ) - facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("missing") assert _run(facade) == "import-1" - update = facade.update_source.call_args.kwargs["payload"] + update = facade.payload("update_source") assert update.status == "active" assert update.metadata["lastImport"]["state"] == "completed" assert update.metadata["pendingImport"] is None - policy = facade.update_source_sync_policy.call_args.kwargs["payload"] + policy = facade.payload("update_source_sync_policy") assert policy.expected_source_version == 6 def test_finalize_source_import_persists_failure_on_visible_source() -> None: - facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace( - failure=SimpleNamespace(message="provider failed"), - id="import-1", - last_error_code="PROVIDER_FAILED", - state="failed", + facade = _facade( + source_workflow_response( + workflow_id="import-1", + state="failed", + failure_message="provider failed", + last_error_code="PROVIDER_FAILED", + ) ) - facade.get_source.return_value = _source() + facade.script("get_source", _source()) + facade.script("update_source", _source(status="error")) assert _run(facade) == "import-1" - update = facade.update_source.call_args.kwargs["payload"] + update = facade.payload("update_source") assert update.status == "error" assert "parameters" not in update.metadata assert update.metadata["lastImport"] == { @@ -192,28 +219,28 @@ def test_finalize_source_import_persists_failure_on_visible_source() -> None: "workflowId": "import-1", } assert update.metadata["pendingImport"] is None - facade.update_source_sync_policy.assert_not_called() + assert facade.call_count("update_source_sync_policy") == 0 def test_finalize_website_source_import_waits_for_policy_write_fence() -> None: - facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace(id="import-1", state="completed") - facade.get_source.return_value = SimpleNamespace( - id="source-1", - metadata={ - "lastImport": { - "kind": "website-crawl-import", - "state": "completed", - "syncPolicy": {"enabled": True, "mode": "interval"}, - "workflowId": "import-1", + facade = _facade(_completed_workflow()) + facade.script( + "get_source", + source_response( + metadata={ + "lastImport": { + "kind": "website-crawl-import", + "state": "completed", + "syncPolicy": {"enabled": True, "mode": "interval"}, + "workflowId": "import-1", + }, + "preview": False, }, - "preview": False, - }, - status="active", - version=5, + status="active", + version=5, + ), ) - facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("missing") - facade.update_source_sync_policy.side_effect = KnowledgeFSProductRequestRejectedError(status_code=400) + facade.script("update_source_sync_policy", KnowledgeFSProductRequestRejectedError(status_code=400)) with pytest.raises(KnowledgeFSSourceImportNotReadyError): _run(facade) @@ -223,25 +250,25 @@ def test_finalize_website_source_import_waits_for_policy_write_fence() -> None: def test_finalize_non_website_source_import_does_not_reclassify_a_bad_request( import_kind: str, ) -> None: - facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace(id="import-1", state="completed") - facade.get_source.return_value = SimpleNamespace( - id="source-1", - metadata={ - "lastImport": { - "kind": import_kind, - "state": "completed", - "syncPolicy": {"enabled": True, "mode": "interval"}, - "workflowId": "import-1", + facade = _facade(_completed_workflow()) + facade.script( + "get_source", + source_response( + metadata={ + "lastImport": { + "kind": import_kind, + "state": "completed", + "syncPolicy": {"enabled": True, "mode": "interval"}, + "workflowId": "import-1", + }, + "preview": False, }, - "preview": False, - }, - status="active", - version=5, + status="active", + version=5, + ), ) - facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("missing") rejection = KnowledgeFSProductRequestRejectedError(status_code=400) - facade.update_source_sync_policy.side_effect = rejection + facade.script("update_source_sync_policy", rejection) with pytest.raises(KnowledgeFSProductRequestRejectedError) as raised: _run(facade) @@ -251,8 +278,8 @@ def test_finalize_non_website_source_import_does_not_reclassify_a_bad_request( def test_finalize_source_import_retries_a_source_revision_conflict(monkeypatch: pytest.MonkeyPatch) -> None: error = KnowledgeFSProductRequestRejectedError(status_code=409) - monkeypatch.setattr(task_module, "finalize_source_import_once", MagicMock(side_effect=error)) - retry = MagicMock(side_effect=Retry()) + monkeypatch.setattr(task_module, "finalize_source_import_once", lambda **_: _raise(error)) + retry = RecordingRetry() monkeypatch.setattr(task_module.finalize_source_import, "retry", retry) with pytest.raises(Retry): @@ -264,7 +291,7 @@ def test_finalize_source_import_retries_a_source_revision_conflict(monkeypatch: workflow_id="import-1", ) - retry.assert_called_once_with(exc=error) + assert retry.calls == [{"exc": error}] @pytest.mark.parametrize("status_code", [400, 422]) @@ -272,8 +299,8 @@ def test_finalize_source_import_does_not_retry_a_permanent_rejection( monkeypatch: pytest.MonkeyPatch, status_code: Literal[400, 422] ) -> None: error = KnowledgeFSProductRequestRejectedError(status_code=status_code) - monkeypatch.setattr(task_module, "finalize_source_import_once", MagicMock(side_effect=error)) - retry = MagicMock(side_effect=Retry()) + monkeypatch.setattr(task_module, "finalize_source_import_once", lambda **_: _raise(error)) + retry = RecordingRetry() monkeypatch.setattr(task_module.finalize_source_import, "retry", retry) with pytest.raises(KnowledgeFSProductRequestRejectedError) as raised: @@ -286,4 +313,4 @@ def test_finalize_source_import_does_not_retry_a_permanent_rejection( ) assert raised.value is error - retry.assert_not_called() + assert retry.calls == [] From 6eee9f7ee58c1a703ea168da3ca07de07f2104e9 Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:49:32 +0800 Subject: [PATCH 5/5] test(api): replace test doubles in KnowledgeFS task tests with typed fakes Co-Authored-By: Claude Opus 5.5 --- ..._knowledge_fs_agent_investigation_tasks.py | 59 ++-- ...est_knowledge_fs_failed_retrieval_tasks.py | 106 ++++--- ...owledge_fs_initial_source_preview_tasks.py | 300 +++++++----------- .../test_knowledge_fs_lifecycle_tasks.py | 80 +++-- .../tasks/test_knowledge_fs_upgrade_tasks.py | 37 ++- 5 files changed, 302 insertions(+), 280 deletions(-) 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 851c0b81a60..cc051fc8853 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 @@ -1,15 +1,30 @@ -from unittest.mock import Mock +from dataclasses import dataclass, field from uuid import uuid4 import pytest from celery.exceptions import Retry from dify_agent.protocol.knowledge_fs import KnowledgeFsError +from dify_agent.protocol.knowledge_investigation import KnowledgeInvestigationPayload 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 +@dataclass +class _Gateway: + """Records investigation captures; raises ``error`` when one is configured.""" + + error: Exception | None = None + captures: list[tuple[KnowledgeInvestigationPayload, str]] = field(default_factory=list) + + def capture_investigation(self, report: KnowledgeInvestigationPayload, control_space_id: str) -> None: + self.captures.append((report, control_space_id)) + if self.error is not None: + raise self.error + + def report(): space = str(uuid4()) return { @@ -23,58 +38,56 @@ def report(): def test_dataset_worker_reconstructs_report_for_each_delivery(monkeypatch): - gateway = Mock() + gateway = _Gateway() monkeypatch.setattr(module, "AgentKnowledgeGateway", lambda: gateway) payload, space = report() task = module.capture_agent_knowledge_investigation_task task.run(report=payload, control_space_id=space) task.run(report=payload, control_space_id=space) assert task_options(task)["queue"] == "dataset" - first, second = gateway.capture_investigation.call_args_list + first, second = gateway.captures assert first == second - assert str(first.args[0].investigation_id) == payload["investigation_id"] + assert first[1] == space + assert str(first[0].investigation_id) == payload["investigation_id"] def test_revoked_authority_stops_publication_and_transient_failures_retry(monkeypatch): - gateway = Mock() + gateway = _Gateway(error=KnowledgeFsError("KNOWLEDGE_SCOPE_CHANGED", "revoked", 403)) monkeypatch.setattr(module, "AgentKnowledgeGateway", lambda: gateway) payload, space = report() task = module.capture_agent_knowledge_investigation_task - retry = Mock(side_effect=Retry()) + retry = RecordingRetry() monkeypatch.setattr(task, "retry", retry) - gateway.capture_investigation.side_effect = KnowledgeFsError("KNOWLEDGE_SCOPE_CHANGED", "revoked", 403) task.run(report=payload, control_space_id=space) - retry.assert_not_called() - gateway.capture_investigation.side_effect = OSError("offline") + assert retry.calls == [] + gateway.error = OSError("offline") with pytest.raises(Retry): task.run(report=payload, control_space_id=space) - retry.assert_called_once() + assert len(retry.calls) == 1 @pytest.mark.parametrize("status_code", [400, 403, 409, 413, 422]) def test_permanent_remote_rejections_do_not_retry(monkeypatch, status_code): - gateway = Mock() - gateway.capture_investigation.side_effect = KnowledgeFSProductRequestRejectedError(status_code=status_code) + gateway = _Gateway(error=KnowledgeFSProductRequestRejectedError(status_code=status_code)) monkeypatch.setattr(module, "AgentKnowledgeGateway", lambda: gateway) task = module.capture_agent_knowledge_investigation_task - retry = Mock(side_effect=Retry()) + retry = RecordingRetry() monkeypatch.setattr(task, "retry", retry) payload, space = report() task.run(report=payload, control_space_id=space) - gateway.capture_investigation.assert_called_once() - retry.assert_not_called() + assert len(gateway.captures) == 1 + assert retry.calls == [] @pytest.mark.parametrize(("retries", "countdown"), [(0, 30), (1, 60), (2, 120)]) def test_rate_limited_report_retries_with_backoff(monkeypatch, retries, countdown): - gateway = Mock() error = KnowledgeFSProductRequestRejectedError(status_code=429) - gateway.capture_investigation.side_effect = error + gateway = _Gateway(error=error) monkeypatch.setattr(module, "AgentKnowledgeGateway", lambda: gateway) task = module.capture_agent_knowledge_investigation_task - retry = Mock(side_effect=Retry()) + retry = RecordingRetry() monkeypatch.setattr(task, "retry", retry) monkeypatch.setattr(task.request, "retries", retries) payload, space = report() @@ -82,17 +95,17 @@ def test_rate_limited_report_retries_with_backoff(monkeypatch, retries, countdow with pytest.raises(Retry): task.run(report=payload, control_space_id=space) - retry.assert_called_once_with(exc=error, countdown=countdown) + assert retry.calls == [{"exc": error, "countdown": countdown}] def test_malformed_queued_report_is_rejected_before_gateway_publication(monkeypatch): - gateway = Mock() + gateway = _Gateway() monkeypatch.setattr(module, "AgentKnowledgeGateway", lambda: gateway) task = module.capture_agent_knowledge_investigation_task - retry = Mock(side_effect=Retry()) + retry = RecordingRetry() monkeypatch.setattr(task, "retry", retry) task.run(report={"investigation_id": "invalid"}, control_space_id=str(uuid4())) - gateway.capture_investigation.assert_not_called() - retry.assert_not_called() + assert gateway.captures == [] + assert retry.calls == [] 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 de3b4cb1dbc..72d167d09bb 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 @@ -1,16 +1,17 @@ from __future__ import annotations -from types import SimpleNamespace +from dataclasses import dataclass, field from typing import Literal, TypedDict -from unittest.mock import MagicMock import pytest from celery.app.task import Context from celery.exceptions import Retry from services.knowledge_fs.app_admission_service import KnowledgeFSAppAdmissionError +from services.knowledge_fs.product_dto import KnowledgeFSWorkflowFailedRetrievalCaptureResponse from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError 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 EVENT_ID = "019fac9f-bfb0-75ee-9af5-252ebafbac1e" @@ -38,14 +39,28 @@ def _task_kwargs() -> FailedRetrievalTaskArguments: } -def _install_failing_capability(monkeypatch: pytest.MonkeyPatch, error: Exception) -> MagicMock: - capability = MagicMock() - capability.capture_workflow_failed_retrieval.side_effect = error - monkeypatch.setattr( - task_module, - "get_knowledge_fs_runtime", - lambda _session_maker: SimpleNamespace(app_capabilities=capability), - ) +@dataclass +class _AppCapabilities: + outcome: KnowledgeFSWorkflowFailedRetrievalCaptureResponse | Exception + calls: list[dict[str, object]] = field(default_factory=list) + + def capture_workflow_failed_retrieval(self, **kwargs: object) -> KnowledgeFSWorkflowFailedRetrievalCaptureResponse: + self.calls.append(kwargs) + if isinstance(self.outcome, Exception): + raise self.outcome + return self.outcome + + +@dataclass(frozen=True) +class _Runtime: + app_capabilities: _AppCapabilities + + +def _install_capability( + monkeypatch: pytest.MonkeyPatch, outcome: KnowledgeFSWorkflowFailedRetrievalCaptureResponse | Exception +) -> _AppCapabilities: + capability = _AppCapabilities(outcome) + monkeypatch.setattr(task_module, "get_knowledge_fs_runtime", lambda _session_maker: _Runtime(capability)) monkeypatch.setattr(task_module.session_factory, "get_session_maker", lambda: object()) return capability @@ -62,23 +77,20 @@ def _run_task_with_retries(retries: int) -> None: def test_failed_retrieval_task_runs_on_dataset_queue_and_reauthorizes_app( monkeypatch: pytest.MonkeyPatch, ) -> None: - capability = MagicMock() - capability.capture_workflow_failed_retrieval.return_value = SimpleNamespace( - failed_query_id="failed-query-1", - verdict="retrieval-miss", - bad_case_id="bad-case-1", + capability = _install_capability( + monkeypatch, + KnowledgeFSWorkflowFailedRetrievalCaptureResponse.model_construct( + failed_query_id="failed-query-1", + verdict="retrieval-miss", + bad_case_id="bad-case-1", + ), ) - monkeypatch.setattr( - task_module, - "get_knowledge_fs_runtime", - lambda _session_maker: SimpleNamespace(app_capabilities=capability), - ) - monkeypatch.setattr(task_module.session_factory, "get_session_maker", lambda: object()) task_module.capture_workflow_failed_retrieval_task.run(**_task_kwargs()) assert task_options(task_module.capture_workflow_failed_retrieval_task)["queue"] == "dataset" - call = capability.capture_workflow_failed_retrieval.call_args.kwargs + assert len(capability.calls) == 1 + call = capability.calls[0] assert call["tenant_id"] == "tenant-1" assert call["app_id"] == "app-1" assert call["resource"].control_space_id == "control-1" @@ -88,67 +100,67 @@ def test_failed_retrieval_task_runs_on_dataset_queue_and_reauthorizes_app( def test_enqueue_is_idempotency_aware_and_never_propagates_broker_failure(monkeypatch: pytest.MonkeyPatch) -> None: - delay = MagicMock() - monkeypatch.setattr(task_module.capture_workflow_failed_retrieval_task, "delay", delay) + delayed: list[dict[str, object]] = [] + monkeypatch.setattr( + task_module.capture_workflow_failed_retrieval_task, "delay", lambda **kwargs: delayed.append(kwargs) + ) task_module.enqueue_workflow_failed_retrieval_capture(**_task_kwargs()) - delay.assert_called_once_with(**_task_kwargs()) + assert delayed == [_task_kwargs()] - delay.side_effect = RuntimeError("broker unavailable") + def unavailable_broker(**kwargs: object) -> None: + delayed.append(kwargs) + raise RuntimeError("broker unavailable") + + monkeypatch.setattr(task_module.capture_workflow_failed_retrieval_task, "delay", unavailable_broker) task_module.enqueue_workflow_failed_retrieval_capture(**_task_kwargs()) + assert len(delayed) == 2 def test_terminal_admission_rejection_is_swallowed(monkeypatch: pytest.MonkeyPatch) -> None: - capability = MagicMock() - capability.capture_workflow_failed_retrieval.side_effect = KnowledgeFSAppAdmissionError("not admitted") - monkeypatch.setattr( - task_module, - "get_knowledge_fs_runtime", - lambda _session_maker: SimpleNamespace(app_capabilities=capability), - ) - monkeypatch.setattr(task_module.session_factory, "get_session_maker", lambda: object()) + capability = _install_capability(monkeypatch, KnowledgeFSAppAdmissionError("not admitted")) task_module.capture_workflow_failed_retrieval_task.run(**_task_kwargs()) - capability.capture_workflow_failed_retrieval.assert_called_once() + assert len(capability.calls) == 1 def test_rate_limit_rejection_retries_with_initial_backoff(monkeypatch: pytest.MonkeyPatch) -> None: error = KnowledgeFSProductRequestRejectedError(status_code=429) - capability = _install_failing_capability(monkeypatch, error) - retry = MagicMock(side_effect=Retry()) + capability = _install_capability(monkeypatch, error) + retry = RecordingRetry() monkeypatch.setattr(task_module.capture_workflow_failed_retrieval_task, "retry", retry) with pytest.raises(Retry): _run_task_with_retries(0) - capability.capture_workflow_failed_retrieval.assert_called_once() - retry.assert_called_once_with(exc=error, countdown=30) + assert len(capability.calls) == 1 + assert retry.calls == [{"exc": error, "countdown": 30}] def test_transient_failure_retries_with_exponential_backoff(monkeypatch: pytest.MonkeyPatch) -> None: error = RuntimeError("KnowledgeFS temporarily unavailable") - capability = _install_failing_capability(monkeypatch, error) - retry = MagicMock(side_effect=Retry()) + capability = _install_capability(monkeypatch, error) + retry = RecordingRetry() monkeypatch.setattr(task_module.capture_workflow_failed_retrieval_task, "retry", retry) with pytest.raises(Retry): _run_task_with_retries(1) - capability.capture_workflow_failed_retrieval.assert_called_once() - retry.assert_called_once_with(exc=error, countdown=60) + assert len(capability.calls) == 1 + assert retry.calls == [{"exc": error, "countdown": 60}] def test_retry_budget_exhaustion_raises_original_failure(monkeypatch: pytest.MonkeyPatch) -> None: error = RuntimeError("KnowledgeFS remained unavailable") - capability = _install_failing_capability(monkeypatch, error) - retry = MagicMock() + capability = _install_capability(monkeypatch, error) + retry = RecordingRetry() monkeypatch.setattr(task_module.capture_workflow_failed_retrieval_task, "retry", retry) with pytest.raises(RuntimeError, match="KnowledgeFS remained unavailable") as exc_info: _run_task_with_retries(3) assert exc_info.value is error - capability.capture_workflow_failed_retrieval.assert_called_once() - retry.assert_not_called() + assert len(capability.calls) == 1 + assert retry.calls == [] 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 a311e611b7d..d98ac2903dd 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 @@ -1,13 +1,58 @@ -from types import SimpleNamespace +from collections.abc import Callable +from dataclasses import dataclass, field from typing import Never -from unittest.mock import MagicMock, call, patch +from unittest.mock import patch -from models.account import AccountStatus +from sqlalchemy import inspect + +from core.db.session_factory import session_factory +from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole from services.knowledge_fs.initial_source_preview import KnowledgeFSInitialSourcePreviewCanceledError -from services.knowledge_fs.product_dto import KnowledgeFSInitialSourcePreviewResponse +from services.knowledge_fs.product_dto import ( + KnowledgeFSInitialSourcePreviewJobResponse, + 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 +_JOB_KEY = {"tenant_id": "tenant-1", "account_id": "account-1", "job_id": "job-1"} + + +@dataclass +class _PreviewJobService: + """Records preview job transitions; ``transition_results`` are consumed in order.""" + + transition_results: list[bool] + job_status: str = "running" + transitions: list[dict[str, object]] = field(default_factory=list) + stored: list[dict[str, object]] = field(default_factory=list) + released: list[dict[str, str]] = field(default_factory=list) + lookups: list[dict[str, str]] = field(default_factory=list) + + def transition_status(self, **kwargs: object) -> bool: + self.transitions.append(kwargs) + return self.transition_results.pop(0) + + def store_content(self, **kwargs: object) -> None: + self.stored.append(kwargs) + + def release_active_job(self, *, tenant_id: str, account_id: str, job_id: str) -> None: + self.released.append({"tenant_id": tenant_id, "account_id": account_id, "job_id": job_id}) + + def get(self, *, tenant_id: str, account_id: str, job_id: str) -> KnowledgeFSInitialSourcePreviewJobResponse: + self.lookups.append({"tenant_id": tenant_id, "account_id": account_id, "job_id": job_id}) + return KnowledgeFSInitialSourcePreviewJobResponse.model_construct(status=self.job_status) + + +@dataclass +class _PreviewService: + run: Callable[..., KnowledgeFSInitialSourcePreviewResponse] + calls: list[dict[str, object]] = field(default_factory=list) + + def preview(self, **kwargs: object) -> KnowledgeFSInitialSourcePreviewResponse: + self.calls.append(kwargs) + return self.run(**kwargs) + def _payload() -> dict[str, object]: return { @@ -20,28 +65,28 @@ def _payload() -> dict[str, object]: } -def test_preview_task_uses_the_dataset_queue() -> None: - assert task_options(run_knowledge_fs_initial_source_preview)["queue"] == "dataset" +def _seed_account(*, member: bool) -> None: + tenant = Tenant(name="Preview tenant") + tenant.id = "tenant-1" + account = Account(name="Preview user", email="preview-user@example.com", status=AccountStatus.ACTIVE) + account.id = "account-1" + with session_factory.create_session() as session: + session.add_all([tenant, account]) + if member: + session.add( + TenantAccountJoin( + tenant_id=tenant.id, + account_id=account.id, + current=True, + role=TenantAccountRole.OWNER, + ) + ) + session.commit() -def test_preview_task_persists_running_and_completed_states() -> None: - account = SimpleNamespace( - current_tenant_id="tenant-1", - id="account-1", - set_tenant_id_with_session=MagicMock(), - status=AccountStatus.ACTIVE, - ) - session = MagicMock() - session.get.return_value = account - session.scalar.return_value = "membership-1" - session_context = MagicMock() - session_context.__enter__.return_value = session - job_service = MagicMock() - job_service.transition_status.side_effect = [True, True] - preview_service = MagicMock() - result = KnowledgeFSInitialSourcePreviewResponse(kind="website_crawl") - preview_service.preview.return_value = result - +def _run_preview(job_service: _PreviewJobService, preview_service: _PreviewService | None) -> list[dict[str, object]]: + """Run the task and return the cleanup schedules it requested.""" + scheduled_cleanups: list[dict[str, object]] = [] with ( patch( "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewJobService", @@ -50,15 +95,11 @@ def test_preview_task_persists_running_and_completed_states() -> None: patch( "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewService", return_value=preview_service, - ), + ) as preview_service_type, patch( - "tasks.knowledge_fs_initial_source_preview_tasks.session_factory.create_session", - return_value=session_context, + "tasks.knowledge_fs_initial_source_preview_tasks.cleanup_knowledge_fs_initial_source_preview.apply_async", + side_effect=lambda **kwargs: scheduled_cleanups.append(kwargs), ), - patch("tasks.knowledge_fs_initial_source_preview_tasks.session_factory.get_session_maker"), - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.cleanup_knowledge_fs_initial_source_preview.apply_async" - ) as schedule_cleanup, ): run_knowledge_fs_initial_source_preview.run( tenant_id="tenant-1", @@ -66,92 +107,53 @@ def test_preview_task_persists_running_and_completed_states() -> None: job_id="job-1", payload=_payload(), ) + if preview_service is None: + preview_service_type.assert_not_called() + return scheduled_cleanups - assert job_service.transition_status.call_args_list == [ - call( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - status="running", - allowed_from=("pending",), - ), - call( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - status="completed", - allowed_from=("running",), - result=result, - ), + +def test_preview_task_uses_the_dataset_queue() -> None: + assert task_options(run_knowledge_fs_initial_source_preview)["queue"] == "dataset" + + +def test_preview_task_persists_running_and_completed_states() -> None: + _seed_account(member=True) + job_service = _PreviewJobService(transition_results=[True, True]) + result = KnowledgeFSInitialSourcePreviewResponse(kind="website_crawl") + preview_service = _PreviewService(run=lambda **_: result) + + scheduled_cleanups = _run_preview(job_service, preview_service) + + assert job_service.transitions == [ + {**_JOB_KEY, "status": "running", "allowed_from": ("pending",)}, + {**_JOB_KEY, "status": "completed", "allowed_from": ("running",), "result": result}, ] - job_service.store_content.assert_called_once_with( - tenant_id="tenant-1", account_id="account-1", job_id="job-1", result=result - ) - schedule_cleanup.assert_called_once_with( - kwargs={"account_id": "account-1", "job_id": "job-1", "tenant_id": "tenant-1"}, countdown=3600 - ) - job_service.release_active_job.assert_called_once_with( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - ) - account.set_tenant_id_with_session.assert_called_once_with("tenant-1", session=session) - session.expunge.assert_called_once_with(account) - preview_service.preview.assert_called_once() + assert job_service.stored == [{**_JOB_KEY, "result": result}] + assert scheduled_cleanups == [ + {"kwargs": {"account_id": "account-1", "job_id": "job-1", "tenant_id": "tenant-1"}, "countdown": 3600} + ] + assert job_service.released == [_JOB_KEY] + assert len(preview_service.calls) == 1 + account = preview_service.calls[0]["account"] + assert isinstance(account, Account) + assert account.id == "account-1" + # The tenant was selected through the task's session before the account was detached from it. + assert account.current_tenant_id == "tenant-1" + assert inspect(account).detached def test_preview_task_does_not_start_a_canceled_job() -> None: - job_service = MagicMock() - job_service.transition_status.return_value = False + job_service = _PreviewJobService(transition_results=[False]) - with ( - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewJobService", - return_value=job_service, - ), - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewService" - ) as preview_service, - patch("tasks.knowledge_fs_initial_source_preview_tasks.session_factory.get_session_maker"), - ): - run_knowledge_fs_initial_source_preview.run( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - payload=_payload(), - ) + _run_preview(job_service, preview_service=None) - job_service.transition_status.assert_called_once_with( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - status="running", - allowed_from=("pending",), - ) - preview_service.assert_not_called() - job_service.release_active_job.assert_called_once_with( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - ) + assert job_service.transitions == [{**_JOB_KEY, "status": "running", "allowed_from": ("pending",)}] + assert job_service.released == [_JOB_KEY] def test_preview_task_keeps_a_cooperatively_canceled_job_terminal() -> None: - account = SimpleNamespace( - current_tenant_id="tenant-1", - id="account-1", - set_tenant_id_with_session=MagicMock(), - status=AccountStatus.ACTIVE, - ) - session = MagicMock() - session.get.return_value = account - session.scalar.return_value = "membership-1" - session_context = MagicMock() - session_context.__enter__.return_value = session - job_service = MagicMock() - job_service.transition_status.return_value = True - job_service.get.return_value = SimpleNamespace(status="canceled") - preview_service = MagicMock() + _seed_account(member=True) + job_service = _PreviewJobService(transition_results=[True], job_status="canceled") def preview(**kwargs: object) -> Never: is_canceled = kwargs["is_canceled"] @@ -159,79 +161,23 @@ def test_preview_task_keeps_a_cooperatively_canceled_job_terminal() -> None: assert is_canceled() is True raise KnowledgeFSInitialSourcePreviewCanceledError("canceled") - preview_service.preview.side_effect = preview - with ( - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewJobService", - return_value=job_service, - ), - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewService", - return_value=preview_service, - ), - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.session_factory.create_session", - return_value=session_context, - ), - patch("tasks.knowledge_fs_initial_source_preview_tasks.session_factory.get_session_maker"), - ): - run_knowledge_fs_initial_source_preview.run( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - payload=_payload(), - ) + preview_service = _PreviewService(run=preview) - job_service.transition_status.assert_called_once_with( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - status="running", - allowed_from=("pending",), - ) + _run_preview(job_service, preview_service) + + assert job_service.transitions == [{**_JOB_KEY, "status": "running", "allowed_from": ("pending",)}] + assert len(preview_service.calls) == 1 + assert job_service.lookups == [_JOB_KEY] def test_preview_task_rejects_an_account_removed_from_the_tenant() -> None: - account = SimpleNamespace( - current_tenant_id=None, - id="account-1", - set_tenant_id_with_session=MagicMock(), - status=AccountStatus.ACTIVE, - ) - session = MagicMock() - session.get.return_value = account - session.scalar.return_value = None - session_context = MagicMock() - session_context.__enter__.return_value = session - job_service = MagicMock() - job_service.transition_status.side_effect = [True, True] + _seed_account(member=False) + job_service = _PreviewJobService(transition_results=[True, True]) - with ( - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewJobService", - return_value=job_service, - ), - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.KnowledgeFSInitialSourcePreviewService" - ) as preview_service, - patch( - "tasks.knowledge_fs_initial_source_preview_tasks.session_factory.create_session", - return_value=session_context, - ), - patch("tasks.knowledge_fs_initial_source_preview_tasks.session_factory.get_session_maker"), - ): - run_knowledge_fs_initial_source_preview.run( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - payload=_payload(), - ) + _run_preview(job_service, preview_service=None) - preview_service.assert_not_called() - assert job_service.transition_status.call_args_list[-1] == call( - tenant_id="tenant-1", - account_id="account-1", - job_id="job-1", - status="failed", - allowed_from=("pending", "running"), - ) + assert job_service.transitions[-1] == { + **_JOB_KEY, + "status": "failed", + "allowed_from": ("pending", "running"), + } 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 b8538d08703..ac63311604d 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 @@ -1,12 +1,14 @@ from __future__ import annotations +from dataclasses import dataclass, field from datetime import timedelta -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from services.knowledge_fs.lifecycle_readiness import KnowledgeFSLifecycleWorkerReadiness +from services.knowledge_fs.lifecycle_saga import KnowledgeFSLifecycleDispatchResult +from services.knowledge_fs.orphan_reconciler import KnowledgeFSOrphanReport from tasks.knowledge_fs_lifecycle_tasks import ( cleanup_knowledge_fs_staged_uploads, run_knowledge_fs_lifecycle_worker, @@ -14,6 +16,41 @@ from tasks.knowledge_fs_lifecycle_tasks import ( from tests.unit_tests.config_override import apply_config_overrides +@dataclass +class _SagaRunner: + results: list[KnowledgeFSLifecycleDispatchResult] + dispatches: list[dict[str, object]] = field(default_factory=list) + + def dispatch_one(self, **kwargs: object) -> KnowledgeFSLifecycleDispatchResult: + self.dispatches.append(kwargs) + return self.results.pop(0) + + +@dataclass +class _OrphanReconciler: + report: KnowledgeFSOrphanReport + calls: list[dict[str, object]] = field(default_factory=list) + + def reconcile(self, *, limit: int, apply_repairs: bool) -> KnowledgeFSOrphanReport: + self.calls.append({"limit": limit, "apply_repairs": apply_repairs}) + return self.report + + +@dataclass +class _StagedUploads: + expired: int + limits: list[int] = field(default_factory=list) + + def cleanup_expired(self, *, limit: int) -> int: + self.limits.append(limit) + return self.expired + + +@dataclass(frozen=True) +class _Runtime: + facade: str + + def test_worker_self_check_returns_before_assembling_remote_when_rollout_is_not_ready() -> None: with ( patch( @@ -31,21 +68,19 @@ def test_worker_self_check_returns_before_assembling_remote_when_rollout_is_not_ def test_ready_worker_dispatches_a_bounded_batch_then_repairs_orphans( monkeypatch: pytest.MonkeyPatch, ) -> None: - runner = SimpleNamespace( - dispatch_one=MagicMock( - side_effect=( - SimpleNamespace(claimed=True, completed=True), - SimpleNamespace(claimed=True, completed=False), - SimpleNamespace(claimed=False, completed=False), - ) - ) + runner = _SagaRunner( + results=[ + KnowledgeFSLifecycleDispatchResult(claimed=True, completed=True, outbox_id="outbox-1"), + KnowledgeFSLifecycleDispatchResult(claimed=True, completed=False, outbox_id="outbox-2"), + KnowledgeFSLifecycleDispatchResult(claimed=False, completed=False, outbox_id=None), + ] ) - reconciler = SimpleNamespace( - reconcile=MagicMock( - return_value=SimpleNamespace( - repaired_control_space_ids=("control-1",), - cleanup_control_space_ids=("control-2", "control-3"), - ) + reconciler = _OrphanReconciler( + report=KnowledgeFSOrphanReport( + repaired_control_space_ids=("control-1",), + remote_orphan_space_ids=(), + local_missing_remote_ids=(), + cleanup_control_space_ids=("control-2", "control-3"), ) ) apply_config_overrides( @@ -69,14 +104,14 @@ def test_ready_worker_dispatches_a_bounded_batch_then_repairs_orphans( result = run_knowledge_fs_lifecycle_worker.run() assert result == {"status": "ok", "dispatched": 2, "completed": 1, "reconciled": 3} - assert runner.dispatch_one.call_count == 3 - assert runner.dispatch_one.call_args.kwargs == { + assert len(runner.dispatches) == 3 + assert runner.dispatches[-1] == { "worker_id": "worker-host:manual", "now": "now-1", "lease_duration": timedelta(seconds=30), "product_enabled": True, } - reconciler.reconcile.assert_called_once_with(limit=5, apply_repairs=True) + assert reconciler.calls == [{"limit": 5, "apply_repairs": True}] def test_staged_upload_cleanup_returns_before_runtime_assembly_when_not_ready() -> None: @@ -94,8 +129,7 @@ def test_staged_upload_cleanup_returns_before_runtime_assembly_when_not_ready() def test_staged_upload_cleanup_uses_the_configured_batch_limit(monkeypatch: pytest.MonkeyPatch) -> None: - staged_uploads = MagicMock() - staged_uploads.cleanup_expired.return_value = 4 + staged_uploads = _StagedUploads(expired=4) apply_config_overrides(monkeypatch, KNOWLEDGE_FS_LIFECYCLE_BATCH_SIZE=25) with ( patch( @@ -108,7 +142,7 @@ def test_staged_upload_cleanup_uses_the_configured_batch_limit(monkeypatch: pyte ), patch( "tasks.knowledge_fs_lifecycle_tasks.get_knowledge_fs_runtime", - return_value=SimpleNamespace(facade="facade-1"), + return_value=_Runtime(facade="facade-1"), ) as runtime, patch( "tasks.knowledge_fs_lifecycle_tasks.KnowledgeFSStagedUploadService", @@ -120,4 +154,4 @@ def test_staged_upload_cleanup_uses_the_configured_batch_limit(monkeypatch: pyte assert result == 4 runtime.assert_called_once_with("maker-1") service.assert_called_once_with("maker-1", facade="facade-1") - staged_uploads.cleanup_expired.assert_called_once_with(limit=25) + assert staged_uploads.limits == [25] 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 47adb1ed5ca..fdc25c17551 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 @@ -1,6 +1,7 @@ from __future__ import annotations -from unittest.mock import MagicMock, patch +from dataclasses import dataclass, field +from unittest.mock import patch import pytest from celery.app.task import Context @@ -15,6 +16,24 @@ from tasks.knowledge_fs_upgrade_tasks import ( from tests.unit_tests.tasks.task_options import task_options +@dataclass +class _UpgradeRunner: + """Returns ``outcome`` from ``run_next`` or raises it, and records failure reports.""" + + outcome: bool | Exception + runs: list[dict[str, object]] = field(default_factory=list) + failures: list[dict[str, object]] = field(default_factory=list) + + def run_next(self, **kwargs: object) -> bool: + self.runs.append(kwargs) + if isinstance(self.outcome, Exception): + raise self.outcome + return self.outcome + + def fail(self, *, job_id: str, error: Exception) -> None: + self.failures.append({"job_id": job_id, "error": error}) + + def test_upgrade_tasks_are_pinned_to_the_dedicated_queue() -> None: assert task_options(run_knowledge_fs_upgrade)["queue"] == KNOWLEDGE_FS_UPGRADE_QUEUE assert task_options(reconcile_knowledge_fs_upgrade_documents)["queue"] == KNOWLEDGE_FS_UPGRADE_QUEUE @@ -33,8 +52,7 @@ def test_deferred_file_cleanup_uses_the_upgrade_session_factory() -> None: @pytest.mark.parametrize("has_more", [True, False]) def test_worker_enqueues_only_the_expected_follow_up(has_more: bool) -> None: - runner = MagicMock() - runner.run_next.return_value = has_more + runner = _UpgradeRunner(outcome=has_more) with ( patch("tasks.knowledge_fs_upgrade_tasks.session_factory.get_session_maker", return_value="maker"), patch("tasks.knowledge_fs_upgrade_tasks.KnowledgeFSUpgradeRunner", return_value=runner), @@ -43,7 +61,8 @@ def test_worker_enqueues_only_the_expected_follow_up(has_more: bool) -> None: ): run_knowledge_fs_upgrade.run(job_id="job-1") - runner.run_next.assert_called_once() + assert len(runner.runs) == 1 + assert runner.runs[0]["job_id"] == "job-1" if has_more: continue_upgrade.assert_called_once_with(kwargs={"job_id": "job-1"}) reconcile.assert_not_called() @@ -53,9 +72,8 @@ def test_worker_enqueues_only_the_expected_follow_up(has_more: bool) -> None: def test_worker_marks_parent_failed_when_a_checkpoint_raises() -> None: - runner = MagicMock() error = RuntimeError("ordinary document upload failed") - runner.run_next.side_effect = error + runner = _UpgradeRunner(outcome=error) with ( patch("tasks.knowledge_fs_upgrade_tasks.session_factory.get_session_maker", return_value="maker"), patch("tasks.knowledge_fs_upgrade_tasks.KnowledgeFSUpgradeRunner", return_value=runner), @@ -63,13 +81,12 @@ def test_worker_marks_parent_failed_when_a_checkpoint_raises() -> None: ): run_knowledge_fs_upgrade.run(job_id="job-1") - runner.fail.assert_called_once_with(job_id="job-1", error=error) + assert runner.failures == [{"job_id": "job-1", "error": error}] def test_worker_marks_parent_failed_when_provisioning_retries_are_exhausted() -> None: - runner = MagicMock() error = KnowledgeFSUpgradeNotReadyError("Space is still provisioning") - runner.run_next.side_effect = error + runner = _UpgradeRunner(outcome=error) run_knowledge_fs_upgrade.request_stack.push(Context(retries=run_knowledge_fs_upgrade.max_retries)) try: with ( @@ -81,4 +98,4 @@ def test_worker_marks_parent_failed_when_provisioning_retries_are_exhausted() -> finally: run_knowledge_fs_upgrade.pop_request() - runner.fail.assert_called_once_with(job_id="job-1", error=error) + assert runner.failures == [{"job_id": "job-1", "error": error}]