RAGcore-Public/tests/integration/indexes/test_alias_locking.py

"""Distributed alias-lock primitive: concurrency, staleness and release."""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from datetime import UTC, datetime, timedelta
from uuid import uuid4

import pytest
from sqlalchemy.ext.asyncio import AsyncEngine

from ragcore.application.indexes.models import IndexOperationStatus
from ragcore.config import Settings
from ragcore.infrastructure.database import create_database_engine, create_session_factory
from ragcore.infrastructure.database.indexes import PostgresIndexOperationStore

pytestmark = pytest.mark.integration


@pytest.fixture
async def engine() -> AsyncIterator[AsyncEngine]:
    engine = create_database_engine(Settings(env="test"))
    try:
        yield engine
    finally:
        await engine.dispose()


@pytest.fixture
def store(engine: AsyncEngine) -> PostgresIndexOperationStore:
    return PostgresIndexOperationStore(create_session_factory(engine))


def alias() -> str:
    return f"rag_dense_lock-test-{uuid4().hex}_active"


@pytest.mark.asyncio
async def test_begin_grants_lock_when_alias_is_free(store: PostgresIndexOperationStore) -> None:
    now = datetime.now(UTC)
    operation_id = await store.begin(
        alias_name=alias(),
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token=uuid4().hex,
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    assert operation_id is not None


@pytest.mark.asyncio
async def test_concurrent_begin_only_one_succeeds(store: PostgresIndexOperationStore) -> None:
    target = alias()
    now = datetime.now(UTC)

    async def attempt(owner: str) -> object:
        return await store.begin(
            alias_name=target,
            operation="build",
            collection_name="rag_dense_x_v1",
            owner_token=owner,
            now=now,
            lease_duration=timedelta(seconds=30),
        )

    results = await asyncio.gather(attempt(uuid4().hex), attempt(uuid4().hex))
    granted = [result for result in results if result is not None]
    assert len(granted) == 1


@pytest.mark.asyncio
async def test_second_promotion_is_rejected_until_first_finishes(
    store: PostgresIndexOperationStore,
) -> None:
    target = alias()
    now = datetime.now(UTC)
    first = await store.begin(
        alias_name=target,
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token="owner-a",
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    assert first is not None

    blocked = await store.begin(
        alias_name=target,
        operation="rollback",
        collection_name="rag_dense_x_v1",
        owner_token="owner-b",
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    assert blocked is None

    released = await store.finish(
        first, "owner-a", IndexOperationStatus.SUCCEEDED, now=now
    )
    assert released is True

    second = await store.begin(
        alias_name=target,
        operation="rollback",
        collection_name="rag_dense_x_v1",
        owner_token="owner-b",
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    assert second is not None


@pytest.mark.asyncio
async def test_expired_lease_is_reclaimed_as_needs_repair(
    store: PostgresIndexOperationStore,
) -> None:
    target = alias()
    started = datetime.now(UTC)
    crashed = await store.begin(
        alias_name=target,
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token="owner-crashed",
        now=started,
        lease_duration=timedelta(seconds=1),
    )
    assert crashed is not None

    later = started + timedelta(seconds=5)
    resumed = await store.begin(
        alias_name=target,
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token="owner-resumed",
        now=later,
        lease_duration=timedelta(seconds=30),
    )
    assert resumed is not None
    assert resumed != crashed

    # The crashed operation must stay detectable, not silently vanish.
    still_blocked_third_party = await store.begin(
        alias_name=target,
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token="owner-third",
        now=later,
        lease_duration=timedelta(seconds=30),
    )
    assert still_blocked_third_party is None


@pytest.mark.asyncio
async def test_finish_is_idempotent_and_requires_terminal_status(
    store: PostgresIndexOperationStore,
) -> None:
    target = alias()
    now = datetime.now(UTC)
    operation_id = await store.begin(
        alias_name=target,
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token="owner-a",
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    assert operation_id is not None

    with pytest.raises(ValueError, match="terminal"):
        await store.finish(operation_id, "owner-a", IndexOperationStatus.IN_PROGRESS, now=now)

    first_finish = await store.finish(
        operation_id, "owner-a", IndexOperationStatus.FAILED_COMPENSATED, now=now
    )
    assert first_finish is True

    second_finish = await store.finish(
        operation_id, "owner-a", IndexOperationStatus.FAILED_COMPENSATED, now=now
    )
    assert second_finish is False


@pytest.mark.asyncio
async def test_lock_scope_is_specific_to_one_alias(store: PostgresIndexOperationStore) -> None:
    now = datetime.now(UTC)
    first_alias, second_alias = alias(), alias()
    first = await store.begin(
        alias_name=first_alias,
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token="owner-a",
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    second = await store.begin(
        alias_name=second_alias,
        operation="build",
        collection_name="rag_dense_y_v1",
        owner_token="owner-b",
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    assert first is not None
    assert second is not None


@pytest.mark.asyncio
async def test_record_step_appends_ordered_entries(store: PostgresIndexOperationStore) -> None:
    target = alias()
    now = datetime.now(UTC)
    operation_id = await store.begin(
        alias_name=target,
        operation="build",
        collection_name="rag_dense_x_v1",
        owner_token="owner-a",
        now=now,
        lease_duration=timedelta(seconds=30),
    )
    assert operation_id is not None

    await store.record_step(operation_id, "owner-a", "created_shadow", now=now)
    await store.record_step(
        operation_id,
        "owner-a",
        "activated",
        now=now,
        details={"previous_collection": "rag_dense_x_v0"},
    )

    # A step recorded by a non-owner must not be applied.
    await store.record_step(operation_id, "owner-mismatched", "hijacked", now=now)