hub-core/tests/test_message_identity.py

from __future__ import annotations

import asyncio
import uuid
from datetime import datetime, timezone

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from hub_core.message_identity import (
    MessageIdentityConflict,
    MessageIdentityError,
    register_message_alias,
)
from hub_core.models.agent_message import AgentMessage
from hub_core.models.message_identity_alias import MessageIdentityAlias
from hub_core.routers.messages import create_messages_router


CANONICAL_ID = uuid.UUID("0b8dd0bf-41d1-47da-96ac-40e443c32e47")
MALFORMED_REFERENCE = "0b8dd0bf-41d-47da-96ac-40e443c32e47"
CREATED_AT = datetime(2026, 8, 20, 6, 9, 1, tzinfo=timezone.utc)


def _database(tmp_path):
    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'messages.db'}")
    sessions = async_sessionmaker(engine, expire_on_commit=False)

    async def prepare() -> None:
        async with engine.begin() as connection:
            await connection.run_sync(AgentMessage.__table__.create)
            await connection.run_sync(MessageIdentityAlias.__table__.create)
        async with sessions() as session:
            session.add(
                AgentMessage(
                    id=CANONICAL_ID,
                    from_agent="risk-nexus",
                    to_agent="the-custodian",
                    subject="Canon request",
                    body="Original body",
                    created_at=CREATED_AT,
                )
            )
            await session.commit()

    asyncio.run(prepare())
    return engine, sessions


def _register(sessions, *, alias=MALFORMED_REFERENCE, message_id=CANONICAL_ID):
    async def run():
        async with sessions() as session:
            return await register_message_alias(
                session,
                alias=alias,
                message_id=message_id,
                source="CUST-IN-0012",
                reason="Historical reference omitted one UUID character",
                created_by="codex",
            )

    return asyncio.run(run())


def _client(sessions) -> TestClient:
    async def get_session():
        async with sessions() as session:
            yield session

    app = FastAPI()
    app.include_router(create_messages_router(get_session))
    return TestClient(app)


def test_alias_registration_is_explicit_attributable_and_idempotent(tmp_path) -> None:
    engine, sessions = _database(tmp_path)
    try:
        first = _register(sessions)
        second = _register(sessions)

        assert first.created is True
        assert second.created is False
        assert first.message_id == second.message_id == CANONICAL_ID
        assert second.source == "CUST-IN-0012"
        assert second.created_by == "codex"
    finally:
        asyncio.run(engine.dispose())


def test_alias_registration_rejects_conflicts_and_canonical_uuid_shadowing(tmp_path) -> None:
    engine, sessions = _database(tmp_path)
    try:
        _register(sessions)
        other_id = uuid.uuid4()

        async def add_other() -> None:
            async with sessions() as session:
                session.add(
                    AgentMessage(
                        id=other_id,
                        from_agent="sender",
                        to_agent="recipient",
                        subject="Other",
                        body="Other body",
                        created_at=CREATED_AT,
                    )
                )
                await session.commit()

        asyncio.run(add_other())
        with pytest.raises(MessageIdentityConflict):
            _register(sessions, message_id=other_id)
        with pytest.raises(MessageIdentityError, match="canonical UUID"):
            _register(sessions, alias=str(other_id), message_id=other_id)
    finally:
        asyncio.run(engine.dispose())


def test_message_operations_resolve_alias_and_return_canonical_identity(tmp_path) -> None:
    engine, sessions = _database(tmp_path)
    try:
        _register(sessions)
        client = _client(sessions)
        before = client.get(f"/messages/thread/{CANONICAL_ID}")
        assert before.status_code == 200
        original_created_at = before.json()[0]["created_at"]

        marked = client.patch(f"/messages/{MALFORMED_REFERENCE}/read")
        assert marked.status_code == 200
        assert marked.json()["id"] == str(CANONICAL_ID)
        assert marked.json()["body"] == "Original body"
        assert marked.json()["created_at"] == original_created_at
        assert marked.json()["read_at"] is not None

        reply = client.post(
            f"/messages/{MALFORMED_REFERENCE}/reply",
            json={"from_agent": "the-custodian", "body": "Handled"},
        )
        assert reply.status_code == 201
        assert reply.json()["thread_id"] == str(CANONICAL_ID)

        thread = client.get(f"/messages/thread/{MALFORMED_REFERENCE}")
        assert thread.status_code == 200
        assert [row["id"] for row in thread.json()][0] == str(CANONICAL_ID)
        assert [row["body"] for row in thread.json()] == ["Original body", "Handled"]

        archived = client.patch(f"/messages/{MALFORMED_REFERENCE}/archive")
        assert archived.status_code == 200
        assert archived.json()["id"] == str(CANONICAL_ID)
        assert archived.json()["archived_at"] is not None
    finally:
        asyncio.run(engine.dispose())


def test_unknown_malformed_reference_is_not_guessed(tmp_path) -> None:
    engine, sessions = _database(tmp_path)
    try:
        client = _client(sessions)
        response = client.patch("/messages/not-a-registered-identity/read")

        assert response.status_code == 404
        assert "not-a-registered-identity" in response.json()["detail"]
    finally:
        asyncio.run(engine.dispose())