roboco/tests/unit/services/test_work_session.py

"""Unit tests for WorkSessionService gateway-backfill methods."""

from __future__ import annotations

import asyncio
from datetime import UTC, datetime
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock
from uuid import UUID, uuid4

import pytest
from roboco.db.tables import AgentTable, ProjectTable, TaskTable, WorkSessionTable
from roboco.models import AgentRole, AgentStatus, Team
from roboco.models.base import Complexity, TaskNature, TaskStatus, TaskType
from roboco.models.work_session import WorkSessionCreate, WorkSessionStatus
from roboco.services.base import ConflictError
from roboco.services.work_session import WorkSessionService, get_work_session_service
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine


def _service() -> WorkSessionService:
    session = MagicMock()
    session.execute = AsyncMock()
    session.flush = AsyncMock()
    return WorkSessionService(session)


def _bind(svc: WorkSessionService, name: str, value: object) -> None:
    object.__setattr__(svc, name, value)


@pytest.mark.asyncio
async def test_files_changed_returns_files_modified_list() -> None:
    svc = _service()
    fake_session = MagicMock(files_modified=["roboco/api/app.py", "README.md"])
    _bind(svc, "get", AsyncMock(return_value=fake_session))
    out = await svc.files_changed(uuid4())
    assert out == ["roboco/api/app.py", "README.md"]


@pytest.mark.asyncio
async def test_files_changed_returns_empty_list_when_session_missing() -> None:
    svc = _service()
    _bind(svc, "get", AsyncMock(return_value=None))
    assert await svc.files_changed(uuid4()) == []


@pytest.mark.asyncio
async def test_files_changed_handles_none_files_modified() -> None:
    svc = _service()
    fake_session = MagicMock(files_modified=None)
    _bind(svc, "get", AsyncMock(return_value=fake_session))
    assert await svc.files_changed(uuid4()) == []


@pytest.mark.asyncio
async def test_has_unpushed_commits_false_when_no_commits() -> None:
    svc = _service()
    fake_session = MagicMock(commits=[], pr_number=None)
    _bind(svc, "get", AsyncMock(return_value=fake_session))
    assert await svc.has_unpushed_commits(uuid4()) is False


@pytest.mark.asyncio
async def test_has_unpushed_commits_true_when_commits_but_no_pr() -> None:
    svc = _service()
    fake_session = MagicMock(commits=["abc", "def"], pr_number=None)
    _bind(svc, "get", AsyncMock(return_value=fake_session))
    assert await svc.has_unpushed_commits(uuid4()) is True


@pytest.mark.asyncio
async def test_has_unpushed_commits_false_when_pr_exists() -> None:
    svc = _service()
    fake_session = MagicMock(commits=["abc"], pr_number=42)
    _bind(svc, "get", AsyncMock(return_value=fake_session))
    assert await svc.has_unpushed_commits(uuid4()) is False


@pytest.mark.asyncio
async def test_has_unpushed_commits_false_when_session_missing() -> None:
    svc = _service()
    _bind(svc, "get", AsyncMock(return_value=None))
    assert await svc.has_unpushed_commits(uuid4()) is False


# ---------------------------------------------------------------------------
# merge_pr must be idempotent + active-guarded like close()/complete()
# ---------------------------------------------------------------------------


def _merge_service(
    ws: object | None = MagicMock(),
) -> tuple[WorkSessionService, MagicMock, AsyncMock]:
    """A service whose session.execute returns ``ws`` via scalar_one_or_none and
    whose flush is observable. Returns ``(svc, session, execute)`` so the tests
    can assert on the lock SELECT and that a guarded no-op never reached flush.
    """
    session = MagicMock()
    execute = AsyncMock()
    if ws is not None:
        result = MagicMock()
        result.scalar_one_or_none.return_value = ws
        execute.return_value = result
    session.execute = execute
    session.flush = AsyncMock()
    svc = WorkSessionService(session)
    return svc, session, execute


@pytest.mark.asyncio
async def test_merge_pr_completes_active_session() -> None:
    """Happy path: an ACTIVE session is recorded as merged + COMPLETED."""

    original_merger = uuid4()
    ws = MagicMock(
        status=WorkSessionStatus.ACTIVE, pr_number=42, branch_name="feature/x"
    )
    svc, session, _execute = _merge_service(ws)

    result = await svc.merge_pr(uuid4(), original_merger)

    assert result is ws
    assert ws.pr_status == "merged"
    assert ws.merged_by == original_merger
    assert ws.pr_merged_at is not None
    assert ws.status == WorkSessionStatus.COMPLETED
    session.flush.assert_awaited_once()


@pytest.mark.asyncio
async def test_merge_pr_idempotent_on_already_completed_preserves_audit_trail() -> None:
    """A retried merge after a successful-but-unconfirmed GitHub merge must NOT
    overwrite the original ``merged_by`` / ``pr_merged_at`` — the merge audit
    trail is preserved. Mirrors close()'s idempotency guard."""

    original_merger = uuid4()
    original_ts = datetime(2026, 6, 1, 12, 0, tzinfo=UTC)
    retry_merger = uuid4()
    ws = MagicMock(
        status=WorkSessionStatus.COMPLETED,
        pr_number=42,
        merged_by=original_merger,
        pr_merged_at=original_ts,
        pr_status="merged",
    )
    svc, session, _execute = _merge_service(ws)

    result = await svc.merge_pr(uuid4(), retry_merger)

    assert result is ws
    # Audit trail untouched — retry did not overwrite the original merger/ts.
    assert ws.merged_by == original_merger
    assert ws.pr_merged_at == original_ts
    assert ws.status == WorkSessionStatus.COMPLETED
    # No mutation happened, so flush must not have run.
    session.flush.assert_not_awaited()


@pytest.mark.asyncio
async def test_merge_pr_does_not_resurrect_abandoned_session() -> None:
    """``merge_pr`` on an ABANDONED session must NOT flip it to COMPLETED — that
    would silently undo the single-active invariant's abandonment and make
    discarded work look like a successful merge."""

    ws = MagicMock(status=WorkSessionStatus.ABANDONED, pr_number=42, merged_by=None)
    svc, session, _execute = _merge_service(ws)

    result = await svc.merge_pr(uuid4(), uuid4())

    assert result is ws
    # Not resurrected — stays abandoned, no merger recorded.
    assert ws.status == WorkSessionStatus.ABANDONED
    assert ws.merged_by is None
    session.flush.assert_not_awaited()


@pytest.mark.asyncio
async def test_merge_pr_locks_work_session_row_with_for_update() -> None:
    """The work_session SELECT must carry FOR UPDATE so concurrent merge_pr
    calls serialize. Mirrors the pr_merge parent-lock assertion in
    test_pr_merge_concurrency.py."""
    ws = MagicMock(
        status=WorkSessionStatus.ACTIVE, pr_number=42, branch_name="feature/x"
    )
    svc, _session, execute = _merge_service(ws)

    await svc.merge_pr(uuid4(), uuid4())

    assert execute.await_args is not None
    stmt = execute.await_args.args[0]
    assert getattr(stmt, "_for_update_arg", None) is not None


# ---------------------------------------------------------------------------
# M37: two concurrent merge_pr calls on the same ACTIVE session must serialize
# — exactly one writes merged_by / pr_merged_at. Exercises real Postgres row
# locking with two separate sessions on asyncio.gather.
# ---------------------------------------------------------------------------


async def _seed_active_session(
    db_session: AsyncSession,
) -> tuple[UUID, UUID, UUID, UUID, UUID]:
    """Seed project + 3 agents + task + ACTIVE work_session; commit and return
    (ws_id, worker_agent_id, merger_a_id, merger_b_id, project_id). The two
    merger agents exist in the agents table so the merged_by FK is satisfied."""
    worker = AgentTable(
        id=uuid4(),
        name="worker",
        slug=f"worker-{uuid4().hex[:8]}",
        role=AgentRole.DEVELOPER,
        team=Team.BACKEND,
        status=AgentStatus.ACTIVE,
        model_config={},
        system_prompt="dev",
        capabilities=[],
        permissions={},
        metrics={},
    )
    merger_a = AgentTable(
        id=uuid4(),
        name="merger-a",
        slug=f"merger-a-{uuid4().hex[:8]}",
        role=AgentRole.CELL_PM,
        team=Team.BACKEND,
        status=AgentStatus.ACTIVE,
        model_config={},
        system_prompt="pm",
        capabilities=[],
        permissions={},
        metrics={},
    )
    merger_b = AgentTable(
        id=uuid4(),
        name="merger-b",
        slug=f"merger-b-{uuid4().hex[:8]}",
        role=AgentRole.CELL_PM,
        team=Team.BACKEND,
        status=AgentStatus.ACTIVE,
        model_config={},
        system_prompt="pm",
        capabilities=[],
        permissions={},
        metrics={},
    )
    for a in (worker, merger_a, merger_b):
        db_session.add(a)
    await db_session.flush()
    project = ProjectTable(
        id=uuid4(),
        name="p",
        slug=f"p-{uuid4().hex[:6]}",
        git_url="git@x:y/z.git",
        assigned_cell=Team.BACKEND,
        created_by=worker.id,
    )
    db_session.add(project)
    await db_session.flush()
    tid = uuid4()
    db_session.add(
        TaskTable(
            id=tid,
            title="t",
            description="d",
            acceptance_criteria=["done"],
            status=TaskStatus.IN_PROGRESS,
            priority=2,
            task_type=TaskType.CODE,
            nature=TaskNature.TECHNICAL,
            estimated_complexity=Complexity.LOW,
            team=Team.BACKEND,
            confirmed_by_human=True,
            project_id=project.id,
            created_by=worker.id,
            branch_name="feature/x",
        )
    )
    await db_session.flush()
    ws_id = uuid4()
    db_session.add(
        WorkSessionTable(
            id=ws_id,
            project_id=project.id,
            task_id=tid,
            agent_id=worker.id,
            branch_name="feature/x",
            base_branch="master",
            target_branch="master",
            status=WorkSessionStatus.ACTIVE,
        )
    )
    await db_session.flush()
    await db_session.commit()
    return (
        ws_id,
        cast("UUID", worker.id),
        cast("UUID", merger_a.id),
        cast("UUID", merger_b.id),
        cast("UUID", project.id),
    )


@pytest.mark.asyncio
async def test_merge_pr_concurrent_calls_serialize_one_write(
    db_session: AsyncSession, _test_database_url: str
) -> None:
    """Two concurrent merge_pr calls on the same ACTIVE session: only one
    writes merged_by / pr_merged_at; the other blocks on the row lock, then
    sees status=COMPLETED and no-ops."""
    (
        ws_id,
        _worker_id,
        merger_a_id,
        merger_b_id,
        _project_id,
    ) = await _seed_active_session(db_session)

    engine = create_async_engine(_test_database_url, future=True)
    factory = async_sessionmaker(
        bind=engine, class_=AsyncSession, expire_on_commit=False
    )
    merger_a = merger_a_id
    merger_b = merger_b_id

    async def _call(merger: UUID) -> tuple[WorkSessionTable | None, UUID]:
        async with factory() as sess:
            svc = get_work_session_service(sess)
            result = await svc.merge_pr(ws_id, merger)
            await sess.commit()
            return result, merger

    # Both fire together: without FOR UPDATE, each session's SELECT sees the
    # last committed (ACTIVE) row and both write their own merger. With FOR
    # UPDATE, one caller's SELECT blocks on the other's row lock until it
    # commits, then reads COMPLETED and no-ops — exactly one merger is
    # recorded. Which caller wins the lock is non-deterministic, so the
    # assertions below check the invariant, not the winner.
    res_a, res_b = await asyncio.gather(_call(merger_a), _call(merger_b))
    await engine.dispose()

    a_row, a_merger = res_a
    b_row, b_merger = res_b

    # Both resolved COMPLETED; both report the same committed merger.
    assert a_row is not None
    assert b_row is not None
    assert a_row.status == b_row.status == WorkSessionStatus.COMPLETED
    winner = a_row.merged_by
    assert winner in (a_merger, b_merger)
    assert b_row.merged_by == winner


# ---------------------------------------------------------------------------
# H10: create() must translate IntegrityError -> ConflictError
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_create_translates_integrity_error_to_conflict(
    db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Two agents racing a same-task active-session create: the 047 partial-
    unique index fires IntegrityError on the loser's flush. The service must
    re-raise as ConflictError so the choreographer's try/except lands a clean
    invalid_state envelope instead of a 500."""
    agent = AgentTable(
        id=uuid4(),
        name="A",
        slug=f"a-{uuid4().hex[:8]}",
        role=AgentRole.DEVELOPER,
        team=Team.BACKEND,
        status=AgentStatus.ACTIVE,
        model_config={},
        system_prompt="dev",
        capabilities=[],
        permissions={},
        metrics={},
    )
    db_session.add(agent)
    await db_session.flush()
    project = ProjectTable(
        id=uuid4(),
        name="p",
        slug=f"p-{uuid4().hex[:6]}",
        git_url="git@x:y/z.git",
        assigned_cell=Team.BACKEND,
        created_by=agent.id,
    )
    db_session.add(project)
    await db_session.flush()
    tid = uuid4()
    db_session.add(
        TaskTable(
            id=tid,
            title="t",
            description="d",
            acceptance_criteria=["done"],
            status=TaskStatus.IN_PROGRESS,
            priority=2,
            task_type=TaskType.CODE,
            nature=TaskNature.TECHNICAL,
            estimated_complexity=Complexity.LOW,
            team=Team.BACKEND,
            confirmed_by_human=True,
            project_id=project.id,
            created_by=agent.id,
            branch_name="feature/x",
        )
    )
    await db_session.flush()

    existing = WorkSessionTable(
        id=uuid4(),
        project_id=project.id,
        task_id=tid,
        agent_id=agent.id,
        branch_name="feature/x",
        base_branch="master",
        target_branch="master",
        status=WorkSessionStatus.ACTIVE,
    )
    db_session.add(existing)
    await db_session.flush()

    svc = get_work_session_service(db_session)

    async def _none(*_a: Any, **_kw: Any) -> None:
        return None

    monkeypatch.setattr(svc, "get_active_for_task_and_agent", _none)

    with pytest.raises(ConflictError):
        await svc.create(
            WorkSessionCreate(
                project_id=cast("UUID", project.id),
                task_id=tid,
                agent_id=cast("UUID", agent.id),
                branch_name="feature/x",
                base_branch="master",
                target_branch="master",
            )
        )