script-build-and-framework-tmp/script_build_host/tests/test_phase_three_ownership.py

from __future__ import annotations

from pathlib import Path

import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker

from script_build_host.domain.errors import (
    MissionAlreadyOwned,
    MissionFencingTokenStale,
    MissionStopRequested,
)
from script_build_host.domain.records import BuildStatus
from script_build_host.infrastructure.legacy_tables import script_build_record
from script_build_host.infrastructure.ownership import FencedCommandGate, OwnerLease
from script_build_host.infrastructure.tables import mission_binding_table
from script_build_host.repositories.legacy_state import SqlAlchemyLegacyBuildStateRepository
from script_build_host.repositories.sqlalchemy import SqlAlchemyMissionBindingRepository


async def _bound_build(sessions: async_sessionmaker[AsyncSession]) -> int:
    state = SqlAlchemyLegacyBuildStateRepository(sessions)
    build_id = await state.create(
        execution_id=1,
        topic_build_id=2,
        topic_id=3,
        agent_type="AigcAgent",
        agent_config={},
        data_source_url=None,
        strategies_config={},
        root_trace_id="root-owner-test",
    )
    await SqlAlchemyMissionBindingRepository(sessions).create(
        script_build_id=build_id,
        root_trace_id="root-owner-test",
        input_snapshot_id=1,
        engine_version="test",
        schema_version="test/v1",
    )
    return build_id


@pytest.mark.asyncio
async def test_owner_lease_excludes_second_owner_and_grows_epoch(
    database: tuple[AsyncEngine, async_sessionmaker[AsyncSession]], tmp_path: Path
) -> None:
    _, sessions = database
    build_id = await _bound_build(sessions)
    first = OwnerLease(sessions, tmp_path / "leases", owner_instance_id="owner-a")
    second = OwnerLease(sessions, tmp_path / "leases", owner_instance_id="owner-b")

    first_token = await first.acquire(build_id)
    with pytest.raises(MissionAlreadyOwned):
        await second.acquire(build_id)
    await first.release()

    second_token = await second.acquire(build_id)
    assert second_token.owner_epoch == first_token.owner_epoch + 1
    with pytest.raises(MissionFencingTokenStale):
        await FencedCommandGate(sessions).verify(first_token)
    await second.release()


@pytest.mark.asyncio
async def test_stop_epoch_fences_current_owner_without_taking_process_lease(
    database: tuple[AsyncEngine, async_sessionmaker[AsyncSession]], tmp_path: Path
) -> None:
    _, sessions = database
    build_id = await _bound_build(sessions)
    lease = OwnerLease(sessions, tmp_path / "leases", owner_instance_id="owner")
    token = await lease.acquire(build_id)
    gate = FencedCommandGate(sessions)
    await gate.verify(token)

    assert await gate.request_stop(build_id) == token.stop_epoch + 1
    with pytest.raises(MissionStopRequested):
        await gate.verify(token)
    assert (
        await SqlAlchemyLegacyBuildStateRepository(sessions).get_status(build_id)
        is BuildStatus.STOPPING
    )
    await lease.release()


@pytest.mark.asyncio
async def test_new_owner_can_cas_stopping_build_to_stopped(
    database: tuple[AsyncEngine, async_sessionmaker[AsyncSession]], tmp_path: Path
) -> None:
    _, sessions = database
    build_id = await _bound_build(sessions)
    old_lease = OwnerLease(sessions, tmp_path / "leases", owner_instance_id="old")
    old_token = await old_lease.acquire(build_id)
    gate = FencedCommandGate(sessions)
    stop_epoch = await gate.request_stop(build_id)
    assert await gate.request_stop(build_id) == stop_epoch
    await old_lease.release()

    takeover = OwnerLease(sessions, tmp_path / "leases", owner_instance_id="takeover")
    token = await takeover.acquire(build_id)
    try:
        assert token.owner_epoch == old_token.owner_epoch + 1
        assert token.stop_epoch == stop_epoch
        await gate.verify_stop_owner(token)
        await gate.complete_stop(token)
        assert (
            await SqlAlchemyLegacyBuildStateRepository(sessions).get_status(build_id)
            is BuildStatus.STOPPED
        )
        with pytest.raises(MissionFencingTokenStale):
            await gate.complete_stop(old_token)
    finally:
        await takeover.release()


@pytest.mark.asyncio
async def test_phase_three_state_and_snapshot_mutations_share_the_owner_fence(
    database: tuple[AsyncEngine, async_sessionmaker[AsyncSession]], tmp_path: Path
) -> None:
    _, sessions = database
    build_id = await _bound_build(sessions)
    lease = OwnerLease(sessions, tmp_path / "leases", owner_instance_id="phase-three-owner")
    token = await lease.acquire(build_id)
    gate = FencedCommandGate(sessions)
    try:
        await gate.compare_and_set_input_snapshot(token, expected_snapshot_id=1, new_snapshot_id=2)
        await gate.begin_phase(token)
        await gate.set_checkpoint(
            token,
            checkpoint_code="PHASE_THREE_ROOT_ACCEPTED",
            summary="accepted",
        )
        await gate.set_checkpoint(
            token,
            checkpoint_code="PHASE_THREE_ROOT_ACCEPTED",
            summary="accepted",
        )
        async with sessions() as session:
            assert (
                await session.scalar(
                    select(mission_binding_table.c.input_snapshot_id).where(
                        mission_binding_table.c.script_build_id == build_id
                    )
                )
                == 2
            )
            row = (
                (
                    await session.execute(
                        select(script_build_record).where(script_build_record.c.id == build_id)
                    )
                )
                .mappings()
                .one()
            )
            assert row["status"] == BuildStatus.PARTIAL.value
            assert row["error_message"] == "PHASE_THREE_ROOT_ACCEPTED"

        await gate.request_stop(build_id)
        with pytest.raises(MissionStopRequested):
            await gate.begin_phase(token)
        with pytest.raises(MissionStopRequested):
            await gate.compare_and_set_input_snapshot(
                token, expected_snapshot_id=2, new_snapshot_id=3
            )
    finally:
        await lease.release()


@pytest.mark.asyncio
async def test_failed_build_cannot_be_resurrected_by_an_owner_token(
    database: tuple[AsyncEngine, async_sessionmaker[AsyncSession]], tmp_path: Path
) -> None:
    _, sessions = database
    build_id = await _bound_build(sessions)
    lease = OwnerLease(sessions, tmp_path / "leases", owner_instance_id="failed-owner")
    token = await lease.acquire(build_id)
    state = SqlAlchemyLegacyBuildStateRepository(sessions)
    await state.set_status(build_id, BuildStatus.FAILED, error_summary="failed")
    gate = FencedCommandGate(sessions)
    try:
        with pytest.raises(MissionFencingTokenStale):
            await gate.begin_phase(token)
        with pytest.raises(MissionFencingTokenStale):
            await gate.set_checkpoint(
                token,
                checkpoint_code="PHASE_THREE_ROOT_ACCEPTED",
                summary="must not restore",
            )
        with pytest.raises(MissionFencingTokenStale):
            await gate.compare_and_set_input_snapshot(
                token, expected_snapshot_id=1, new_snapshot_id=2
            )
        assert await state.get_status(build_id) is BuildStatus.FAILED
    finally:
        await lease.release()