next-tjwater-server/tests/unit/test_task_questions.py

from copy import deepcopy
from datetime import datetime, timedelta, timezone
from uuid import UUID, uuid4
from unittest.mock import AsyncMock

import pytest

from app.auth.permissions import OPERATIONAL_TASK_EXECUTE, OPERATIONAL_TASK_WRITE
from app.domain.domain_errors import DomainError
from app.domain.schemas.operational_task import AgentRunCommand, AgentRunFailureCommand, OperationalTaskRetryCommand
from app.domain.schemas.task_question import TaskQuestionCreate, TaskQuestionReply
from app.services.operational_tasks import OperationalTaskService
from app.services.task_questions import create_question, answer_question
from tests.support.operational_task_repository import TaskRepository, _context, _task_document, _active_run


@pytest.fixture
def anyio_backend():
    return "asyncio"


def setup():
    now = datetime(2026, 10, 30, tzinfo=timezone.utc)
    context = _context()
    task = _task_document(context, now, lease=now + timedelta(minutes=1))
    repo = TaskRepository(task, _active_run(task, now))
    service = OperationalTaskService(repo, context, {OPERATIONAL_TASK_EXECUTE, OPERATIONAL_TASK_WRITE}, clock=lambda: now)
    command = TaskQuestionCreate(request_id=uuid4(), source_session_id="session-1", run_id=repo.run["run_id"],
        questions=[dict(question_id="flow", title="漏水量依据", body="是否有计量记录？", options=[{"id": "unknown", "label": "目前未知"}],
                        mode="single", allow_custom=True, required=True)])
    return service, repo, UUID(task["task_id"]), command


async def finish(service, repo, task_id):
    return await service.complete_run(task_id, UUID(repo.run["run_id"]),
        AgentRunCommand(request_id=uuid4(), expected_revision=repo.run["revision"], source_session_id="session-1"))


def reply(**overrides):
    return TaskQuestionReply(request_id=uuid4(), expected_version=1,
        answers=[{"question_id": "flow", "selected_option_ids": ["unknown"], "custom_text": "需要采用情景假设"}], **overrides)


@pytest.mark.anyio
async def test_question_waits_without_final_report_and_answer_requeues_once_without_fabricating_evidence():
    service, repo, task_id, command = setup()
    question = await create_question(service, task_id, command)
    assert (await create_question(service, task_id, command)).question_request_id == question.question_request_id
    assert len(repo.notifications) == 1
    with pytest.raises(DomainError, match="waiting_human"):
        await answer_question(service, task_id, question.question_request_id, reply())
    waiting = await finish(service, repo, task_id)
    assert waiting.lifecycle_status == "waiting_human" and waiting.active_run is None
    assert repo.run["status"] == "succeeded" and repo.run["failure_code"] is None
    assert len(repo.notifications) == 1 and repo.notifications[0]["subject_type"] == "question"
    answer = reply()
    resumed = await answer_question(service, task_id, question.question_request_id, answer)
    replay = await answer_question(service, task_id, question.question_request_id, answer)
    assert resumed.execution_generation == replay.execution_generation == 2
    assert resumed.active_run.status == "queued"
    assert resumed.questions[0].answers[0].custom_text == "需要采用情景假设"
    assert resumed.questions[0].answered_by == service._actor()
    assert not repo.evidence and not resumed.actions
    assert repo.notifications[0]["lifecycle_status"] == "resolved"


@pytest.mark.anyio
@pytest.mark.parametrize("answers", [[],
    [{"question_id": "foreign", "selected_option_ids": [], "custom_text": "x"}],
    [{"question_id": "flow", "selected_option_ids": ["bogus"], "custom_text": ""}],
    [{"question_id": "flow", "selected_option_ids": ["unknown", "unknown"], "custom_text": ""}],
])
async def test_invalid_answers_do_not_resume_or_resolve_notification(answers):
    service, repo, task_id, command = setup()
    question = await create_question(service, task_id, command)
    await finish(service, repo, task_id)
    with pytest.raises(DomainError) as caught:
        await answer_question(service, task_id, question.question_request_id,
            TaskQuestionReply(request_id=uuid4(), expected_version=1, answers=answers))
    assert caught.value.code == "ANSWER_INVALID"
    assert repo.task["execution_generation"] == 1
    assert repo.notifications[0]["lifecycle_status"] == "active"


@pytest.mark.anyio
async def test_permissions_foreign_session_duplicate_pending_and_stale_generation():
    service, repo, task_id, command = setup()
    service.permissions = frozenset()
    with pytest.raises(DomainError) as caught:
        await create_question(service, task_id, command)
    assert caught.value.code == "PERMISSION_DENIED"
    service.permissions = {OPERATIONAL_TASK_EXECUTE, OPERATIONAL_TASK_WRITE}
    with pytest.raises(DomainError) as caught:
        await create_question(service, task_id, command.model_copy(update={"source_session_id": "foreign"}))
    assert caught.value.code == "SESSION_MISMATCH"
    question = await create_question(service, task_id, command)
    with pytest.raises(DomainError) as caught:
        await create_question(service, task_id, command.model_copy(update={"request_id": uuid4()}))
    assert caught.value.code == "QUESTION_PENDING"
    await finish(service, repo, task_id)
    await service.retry_task(task_id, OperationalTaskRetryCommand(request_id=uuid4(), reason="重新调查"))
    with pytest.raises(DomainError) as caught:
        await answer_question(service, task_id, question.question_request_id, reply())
    assert caught.value.code == "VERSION_CONFLICT"
    assert repo.questions[str(question.question_request_id)]["status"] == "cancelled"
    assert repo.notifications[0]["lifecycle_status"] == "resolved"


@pytest.mark.anyio
async def test_pending_calculations_must_finish_before_asking():
    service, repo, task_id, command = setup()
    async def pending(*_args):
        return [str(uuid4())]
    repo.pending_calculation_jobs = pending
    with pytest.raises(DomainError) as caught:
        await create_question(service, task_id, command)
    assert caught.value.code == "PENDING_CALCULATIONS"
    assert not repo.questions and not repo.notifications


@pytest.mark.anyio
async def test_answer_requires_write_permission_and_matching_version():
    service, repo, task_id, command = setup()
    question = await create_question(service, task_id, command)
    await finish(service, repo, task_id)
    service.permissions = {OPERATIONAL_TASK_EXECUTE}
    with pytest.raises(DomainError) as caught:
        await answer_question(service, task_id, question.question_request_id, reply())
    assert caught.value.code == "PERMISSION_DENIED"
    service.permissions = {OPERATIONAL_TASK_WRITE}
    with pytest.raises(DomainError) as caught:
        await answer_question(service, task_id, question.question_request_id, reply().model_copy(update={"expected_version": 9}))
    assert caught.value.code == "VERSION_CONFLICT"
    assert repo.task["execution_generation"] == 1


@pytest.mark.anyio
async def test_real_execution_failure_cancels_question_instead_of_disguising_failure():
    service, repo, task_id, command = setup()
    question = await create_question(service, task_id, command)
    result = await service.fail_run(task_id, UUID(repo.run["run_id"]), AgentRunFailureCommand(
        request_id=uuid4(), expected_revision=2, source_session_id="session-1",
        failure_code="SOURCE_UNAVAILABLE", failure_detail="Required tool failed"))
    assert result.lifecycle_status == "failed"
    assert repo.questions[str(question.question_request_id)]["status"] == "cancelled"
    assert repo.notifications[0]["lifecycle_status"] == "resolved"


@pytest.mark.anyio
@pytest.mark.parametrize("timeout_count", [1, 3])
async def test_lease_timeout_cancels_old_question_and_unblocks_the_next_generation(timeout_count):
    service, repo, task_id, command = setup()
    question = await create_question(service, task_id, command)
    service.clock = lambda: question.created_at + timedelta(minutes=2)
    repo.claim_expired_runs = AsyncMock(return_value=[deepcopy(repo.run)])
    repo.consecutive_timed_out_runs = AsyncMock(return_value=timeout_count)

    await service.recover_expired_leases(1)

    assert repo.questions[str(question.question_request_id)]["status"] == "cancelled"
    assert repo.notifications[0]["lifecycle_status"] == "resolved"
    with pytest.raises(DomainError):
        await answer_question(service, task_id, question.question_request_id, reply())
    if timeout_count == 3:
        assert repo.task["lifecycle_status"] == "failed"
        return

    await service.claim_run(task_id, UUID(repo.run["run_id"]), AgentRunCommand(
        request_id=UUID(repo.run["execution_request_id"]), expected_revision=1,
        source_session_id="session-2"))
    next_question = await create_question(service, task_id, command.model_copy(update={
        "request_id": uuid4(), "run_id": UUID(repo.run["run_id"]), "source_session_id": "session-2"}))
    assert next_question.status == "pending"
    assert next_question.execution_generation == 2


@pytest.mark.anyio
async def test_multiple_selection_and_optional_question_are_preserved():
    service, repo, task_id, command = setup()
    question_data = command.questions[0].model_dump()
    question_data.update(mode="multiple", allow_custom=False,
        options=[{"id": "in", "label": "入口计量"}, {"id": "out", "label": "出口计量"}])
    command = TaskQuestionCreate(**{**command.model_dump(), "questions": [question_data,
        {**question_data, "question_id": "optional", "required": False}]})
    question = await create_question(service, task_id, command)
    await finish(service, repo, task_id)
    result = await answer_question(service, task_id, question.question_request_id, TaskQuestionReply(
        request_id=uuid4(), expected_version=1,
        answers=[{"question_id": "flow", "selected_option_ids": ["in", "out"], "custom_text": ""}]))
    assert result.execution_generation == 2
    assert result.questions[0].answers[0].selected_option_ids == ["in", "out"]