deerflow-code/offline-backend-20260512/backend/tests/test_deep_research_start_job.py

"""Regression coverage for idempotent Deep Research job submission."""

from types import SimpleNamespace

import pytest

from app.gateway.routers import deep_research


class _SessionStore:
    def __init__(self) -> None:
        self.updates: list[dict] = []

    async def get(self, session_id: str, *, user_id: str):  # noqa: ARG002
        return {
            "id": session_id,
            "query": "测试主题",
            "title": "测试主题",
            "config_snapshot": {"mode": "basic"},
        }

    async def count_active_by_user(self, *, user_id: str) -> int:  # noqa: ARG002
        return 0

    async def update(self, session_id: str, *, user_id: str, **fields):  # noqa: ARG002
        self.updates.append(fields)
        return {"id": session_id, **fields}


class _JobStore:
    async def get_active_for_session(self, session_id: str, *, user_id: str):  # noqa: ARG002
        return None

    async def try_create_or_get_active(self, **kwargs):  # noqa: ANN003
        # Simulate a retry after an earlier request committed the job but did
        # not finish responding to the browser.
        return {"id": "drj_existing", "status": "queued"}, False


class _Dispatcher:
    def __init__(self) -> None:
        self.nudges = 0

    def nudge(self) -> None:
        self.nudges += 1


def _request(*, sessions, jobs, dispatcher):
    return SimpleNamespace(
        app=SimpleNamespace(
            state=SimpleNamespace(
                deep_research_session_store=sessions,
                deep_research_job_store=jobs,
                deep_research_dispatcher=dispatcher,
            )
        )
    )


@pytest.mark.asyncio
async def test_retry_of_existing_queued_job_nudges_dispatcher(monkeypatch):
    """An idempotent retry must recover a queued-but-not-yet-dispatched job."""

    async def fake_require_user(request):  # noqa: ANN001
        return "u1"

    monkeypatch.setattr(deep_research, "_require_user", fake_require_user)
    sessions = _SessionStore()
    dispatcher = _Dispatcher()
    request = _request(sessions=sessions, jobs=_JobStore(), dispatcher=dispatcher)

    response = await deep_research.start_job(
        "drs_1",
        deep_research.StartJobRequest(entry="legacy"),
        request,
    )

    assert response == {
        "jobId": "drj_existing",
        "sessionId": "drs_1",
        "status": "queued",
        "reused": True,
    }
    assert dispatcher.nudges == 1
    assert sessions.updates == [{"status": "running", "active_job_id": "drj_existing"}]


@pytest.mark.asyncio
async def test_queued_job_still_starts_when_session_projection_temporarily_fails(monkeypatch):
    """A committed job is more important than its non-authoritative UI projection."""

    async def fake_require_user(request):  # noqa: ANN001
        return "u1"

    class _ProjectionFailureSessionStore(_SessionStore):
        async def update(self, session_id: str, *, user_id: str, **fields):  # noqa: ARG002
            raise RuntimeError("temporary database readback failure")

    class _NewJobStore(_JobStore):
        async def try_create_or_get_active(self, **kwargs):  # noqa: ANN003
            return {"id": "drj_new", "status": "queued"}, True

    monkeypatch.setattr(deep_research, "_require_user", fake_require_user)
    dispatcher = _Dispatcher()
    response = await deep_research.start_job(
        "drs_1",
        deep_research.StartJobRequest(entry="legacy"),
        _request(
            sessions=_ProjectionFailureSessionStore(),
            jobs=_NewJobStore(),
            dispatcher=dispatcher,
        ),
    )

    assert response["jobId"] == "drj_new"
    assert response["reused"] is False
    assert dispatcher.nudges == 1


@pytest.mark.asyncio
async def test_chat_report_card_can_enable_images_after_collection(monkeypatch):
    """The write-time card override must win over the collection-time default."""

    async def fake_require_user(request):  # noqa: ANN001
        return "u1"

    class _ChatSessionStore(_SessionStore):
        async def get(self, session_id: str, *, user_id: str):  # noqa: ARG002
            return {
                "id": session_id,
                "query": "测试主题",
                "title": "测试主题",
                "runtime_thread_id": "collector-thread",
                "config_snapshot": {
                    "mode": "basic",
                    "collection_mode": "chat",
                    "generate_images": False,
                    "max_generated_images": 2,
                },
            }

    class _CapturingJobStore(_JobStore):
        def __init__(self) -> None:
            self.created: dict | None = None

        async def try_create_or_get_active(self, **kwargs):  # noqa: ANN003
            self.created = kwargs
            return {"id": "drj_images", "status": "queued"}, True

    monkeypatch.setattr(deep_research, "_require_user", fake_require_user)
    sessions = _ChatSessionStore()
    jobs = _CapturingJobStore()
    response = await deep_research.start_job(
        "drs_1",
        deep_research.StartJobRequest(
            entry="chat",
            config={"generate_images": True, "max_generated_images": 3},
        ),
        _request(sessions=sessions, jobs=jobs, dispatcher=_Dispatcher()),
    )

    assert response["jobId"] == "drj_images"
    assert jobs.created is not None
    snapshot = jobs.created["input_snapshot"]["config"]
    assert snapshot["generate_images"] is True
    assert snapshot["max_generated_images"] == 3
    # The session snapshot is persisted too, so a reload/replay keeps the
    # write-time choice instead of reverting to the creation-time false.
    assert any(
        update.get("config_snapshot", {}).get("generate_images") is True
        for update in sessions.updates
    )


@pytest.mark.asyncio
async def test_stream_stays_open_when_durable_replay_has_a_transient_failure(monkeypatch):
    """A stream DB blip must yield a recoverable warning, not an ASGI error group."""

    async def fake_require_user(request):  # noqa: ANN001
        return "u1"

    class _StreamJobStore:
        async def get(self, job_id: str, *, user_id: str):  # noqa: ARG002
            return {"id": job_id, "session_id": "drs_1"}

        async def get_unscoped(self, job_id: str):  # noqa: ARG002
            return None

    class _UnavailableEventStore:
        async def list_after(self, *args, **kwargs):  # noqa: ANN002, ANN003
            raise RuntimeError("temporary replica failure")

    async def is_connected() -> bool:
        return False

    monkeypatch.setattr(deep_research, "_require_user", fake_require_user)
    request = SimpleNamespace(
        app=SimpleNamespace(
            state=SimpleNamespace(
                deep_research_job_store=_StreamJobStore(),
                deep_research_event_store=_UnavailableEventStore(),
            )
        ),
        is_disconnected=is_connected,
    )

    response = await deep_research.stream_job("drj_1", request)
    chunk = await anext(response.body_iterator)
    await response.body_iterator.aclose()

    assert "durable_replay_temporarily_unavailable" in chunk