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.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.ownership import FencedCommandGate, OwnerLease
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()