Merge branch 'claude/guards-tasks-c80f5b' into deploy/konwledge

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:

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

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

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"),

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"))

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)

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(

api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py

@@ -0,0 +1,227 @@
+"""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 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
+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)
+
+
+@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)

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 == []

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 == []

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"),
+    }

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

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]

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 == []

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}]

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),