omnigent-ai--omnigent/tests/db/test_migration_agents_session_id.py

"""Tests for the agents schema migration that replaces session_id with kind."""

from __future__ import annotations

from collections.abc import Iterator
from pathlib import Path

import pytest
import sqlalchemy as sa
from alembic import command
from sqlalchemy.engine import Engine

from omnigent.db.utils import (
    _build_alembic_config,
    clear_engine_cache,
    get_or_create_engine,
)


@pytest.fixture
def db_engine(tmp_path) -> Iterator[Engine]:
    """Fresh SQLite database with the full migration chain applied."""
    db_path = tmp_path / "test.db"
    uri = f"sqlite:///{db_path}"
    engine = get_or_create_engine(uri)
    try:
        yield engine
    finally:
        clear_engine_cache()


def test_agents_kind_column_exists_and_is_not_nullable(db_engine: Engine) -> None:
    """agents.kind is a NOT NULL column added by the migration."""
    columns = {c["name"]: c for c in sa.inspect(db_engine).get_columns("agents")}
    assert "kind" in columns, "agents.kind column must exist after migration"
    assert not columns["kind"]["nullable"], "agents.kind must be NOT NULL"


def test_agents_session_id_column_removed(db_engine: Engine) -> None:
    """agents.session_id must no longer exist after the migration."""
    columns = {c["name"] for c in sa.inspect(db_engine).get_columns("agents")}
    assert "session_id" not in columns, "agents.session_id must be dropped by migration"


def test_agents_session_id_index_removed(db_engine: Engine) -> None:
    """ix_agents_session_id must no longer exist after the migration."""
    index_names = {i["name"] for i in sa.inspect(db_engine).get_indexes("agents")}
    assert "ix_agents_session_id" not in index_names


def test_ix_conversations_agent_id_added(db_engine: Engine) -> None:
    """ix_conversations_agent_id must be present after the migration."""
    index_names = {i["name"] for i in sa.inspect(db_engine).get_indexes("conversations")}
    assert "ix_conversations_agent_id" in index_names


def test_agents_name_index_exists(db_engine: Engine) -> None:
    """The template-name partial unique index is replaced by a plain name index.

    Template-name uniqueness now lives in the store (MySQL has no partial
    indexes); the DB keeps only a non-unique lookup index on
    ``(workspace_id, name, kind, id)`` — kind is included so the template
    lookup seeks past same-named session copies.
    """
    indexes = {i["name"]: i for i in sa.inspect(db_engine).get_indexes("agents")}
    assert "ix_agents_template_name" not in indexes
    assert "ix_agents_name" in indexes
    assert not indexes["ix_agents_name"]["unique"]
    assert indexes["ix_agents_name"]["column_names"] == ["workspace_id", "name", "kind", "id"]


def test_template_agent_kind_stored_and_read(db_engine: Engine) -> None:
    """A template agent inserted with kind='template' round-trips correctly."""
    with db_engine.begin() as conn:
        conn.execute(
            sa.text(
                # kind=1 → 'template'
                "INSERT INTO agents (id, created_at, name, bundle_location, version, kind)"
                " VALUES (:id, :ts, :name, :loc, 1, 1)"
            ),
            {"id": "ag_tmpl", "ts": 1700000001, "name": "my-template", "loc": "ag_tmpl/bundle"},
        )
        kind = conn.execute(
            sa.text("SELECT kind FROM agents WHERE id = :id"), {"id": "ag_tmpl"}
        ).scalar_one()
    assert kind == 1


def test_session_agent_kind_stored_and_read(db_engine: Engine) -> None:
    """A session-scoped agent inserted with kind='session' round-trips correctly."""
    with db_engine.begin() as conn:
        conn.execute(
            sa.text(
                # kind=2 → 'session'
                "INSERT INTO agents (id, created_at, name, bundle_location, version, kind)"
                " VALUES (:id, :ts, :name, :loc, 1, 2)"
            ),
            {"id": "ag_sess", "ts": 1700000001, "name": "my-session", "loc": "ag_sess/bundle"},
        )
        kind = conn.execute(
            sa.text("SELECT kind FROM agents WHERE id = :id"), {"id": "ag_sess"}
        ).scalar_one()
    assert kind == 2


def test_agents_session_id_fk_accepts_existing_session(db_engine: Engine) -> None:
    """conversations.agent_id (forward pointer) accepts a valid agent id."""
    with db_engine.begin() as conn:
        conn.execute(
            sa.text(
                # kind=2 → 'session'
                "INSERT INTO agents (id, created_at, name, bundle_location, version, kind)"
                " VALUES (:id, :ts, :name, :loc, 1, 2)"
            ),
            {"id": "ag_bound", "ts": 1700000001, "name": "bound-agent", "loc": "ag_bound/bundle"},
        )
        conn.execute(
            sa.text(
                "INSERT INTO conversations"
                " (id, created_at, updated_at, root_conversation_id, kind, agent_id)"
                " VALUES (:id, :ts, :ts, :id, 1, :agent_id)"
            ),
            {"id": "conv_bound", "ts": 1700000002, "agent_id": "ag_bound"},
        )
        stored = conn.execute(
            sa.text("SELECT agent_id FROM conversations WHERE id = :id"),
            {"id": "conv_bound"},
        ).scalar_one()
    assert stored == "ag_bound"


def test_agents_session_id_fk_rejects_missing_session(db_engine: Engine) -> None:
    """Without DB FK, conversations.agent_id accepts any value including nonexistent agents.

    Referential integrity is now the application's responsibility.
    """
    # No IntegrityError expected — FK has been removed.
    with db_engine.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO conversations"
                " (id, created_at, updated_at, root_conversation_id, kind, agent_id)"
                " VALUES (:id, :ts, :ts, :id, 1, :agent_id)"
            ),
            {"id": "conv_missing", "ts": 1700000002, "agent_id": "ag_nonexistent"},
        )
    # Clean up
    with db_engine.begin() as conn:
        conn.execute(sa.text("DELETE FROM conversations WHERE id = 'conv_missing'"))


def test_agents_allow_duplicate_template_names_at_db_layer(
    db_engine: Engine,
) -> None:
    """The DB no longer rejects duplicate template names.

    Uniqueness moved from a partial unique index to the store layer (MySQL
    has no partial indexes), so a raw double-insert succeeds; the guard lives
    in ``SqlAlchemyAgentStore.create`` (see tests/stores/test_agent_store.py).
    """
    with db_engine.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO agents (id, created_at, name, bundle_location, version, kind)"
                " VALUES (:id1, :ts, 'dup-template', :loc1, 1, 1),"
                "        (:id2, :ts, 'dup-template', :loc2, 1, 1)"
            ),
            {
                "id1": "ag_dup1",
                "id2": "ag_dup2",
                "ts": 1700000001,
                "loc1": "ag_dup1/bundle",
                "loc2": "ag_dup2/bundle",
            },
        )
        count = conn.execute(
            sa.text("SELECT COUNT(*) FROM agents WHERE name = 'dup-template'")
        ).scalar_one()
        assert count == 2
        conn.execute(sa.text("DELETE FROM agents WHERE name = 'dup-template'"))


def test_agents_session_id_allows_duplicate_names_for_distinct_sessions(
    db_engine: Engine,
) -> None:
    """Two session-scoped agent copies can reuse the same spec name."""
    with db_engine.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO agents (id, created_at, name, bundle_location, version, kind)"
                " VALUES (:id1, :ts, 'shared-name', :loc1, 1, 2),"
                "        (:id2, :ts, 'shared-name', :loc2, 1, 2)"
            ),
            {
                "id1": "ag_s1",
                "id2": "ag_s2",
                "ts": 1700000001,
                "loc1": "ag_s1/bundle",
                "loc2": "ag_s2/bundle",
            },
        )
        count = conn.execute(
            sa.text("SELECT COUNT(*) FROM agents WHERE name = 'shared-name'")
        ).scalar_one()
    assert count == 2


def test_upgrade_does_not_cascade_delete_conversations(tmp_path: Path) -> None:
    """Upgrade must not cascade-delete conversations bound to session-scoped agents.

    On SQLite, batch_alter_table drops and recreates the agents table.  If
    PRAGMA foreign_keys is ON, the DROP fires the ON DELETE CASCADE on
    conversations.agent_id → agents.id and silently wipes every conversation
    that owns an agent.  This test asserts that upgrade preserves them.
    """
    db_path = tmp_path / "upgrade_cascade.db"
    uri = f"sqlite:///{db_path}"

    # Build a raw engine (no auto-migration) to set up the pre-our-migration state.
    raw_engine = sa.create_engine(uri)

    # Migrate to the revision just before ours.
    config = _build_alembic_config(uri)
    with raw_engine.begin() as conn:
        config.attributes["connection"] = conn
        command.upgrade(config, "n1a2b3c4d5e6")

    # Seed one template agent, one session-scoped agent, and the conversation
    # bound to it — exactly the data that would be wiped by the cascade bug.
    with raw_engine.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO agents (id, created_at, name, bundle_location, version)"
                " VALUES ('ag_tmpl', 1, 'my-template', 'ag_tmpl/b', 1),"
                "        ('ag_sess', 2, 'my-session', 'ag_sess/b', 1)"
            )
        )
        conn.execute(
            sa.text(
                # Seeded at n1 and upgraded only to o1 — both precede the enum→
                # SMALLINT migration (q1), so conversations.kind is still a string.
                "INSERT INTO conversations"
                " (id, created_at, updated_at, root_conversation_id, kind, agent_id)"
                " VALUES ('conv_1', 3, 3, 'conv_1', 'default', 'ag_sess')"
            )
        )
        conn.execute(sa.text("UPDATE agents SET session_id = 'conv_1' WHERE id = 'ag_sess'"))

    # Run our migration.
    config2 = _build_alembic_config(uri)
    with raw_engine.begin() as conn:
        config2.attributes["connection"] = conn
        command.upgrade(config2, "o1a2b3c4d5e6")

    with raw_engine.begin() as conn:
        conv_ids = list(conn.execute(sa.text("SELECT id FROM conversations")).scalars())
        agent_kinds = {
            row[0]: row[1] for row in conn.execute(sa.text("SELECT id, kind FROM agents"))
        }

    assert "conv_1" in conv_ids, "Upgrade must not cascade-delete bound conversations"
    assert agent_kinds.get("ag_sess") == "session"
    assert agent_kinds.get("ag_tmpl") == "template"

    raw_engine.dispose()
    clear_engine_cache()


def test_agents_session_id_downgrade_round_trip(tmp_path: Path) -> None:
    """Downgrade restores session_id from conversations.agent_id and drops kind.

    Uses a raw engine (no auto-migration) to avoid SQLite FK enforcement issues:
    the o1a2b3c4d5e6 downgrade re-adds fk_agents_session_id (ON DELETE CASCADE),
    and subsequent batch_alter_table calls on conversations would cascade-delete
    agents if PRAGMA foreign_keys is ON. A raw engine keeps FK enforcement off.
    """
    db_path = tmp_path / "downgrade.db"
    uri = f"sqlite:///{db_path}"

    # Use a raw engine (no auto-migration) so PRAGMA foreign_keys stays OFF,
    # avoiding cascade issues from the re-added fk_agents_session_id FK.
    raw_engine = sa.create_engine(uri)

    # Migrate to current head first.
    config = _build_alembic_config(uri)
    with raw_engine.begin() as conn:
        config.attributes["connection"] = conn
        command.upgrade(config, "head")

    # Seed data on the upgraded schema: one template, one session-scoped agent.
    # This runs against the full chain (head), where agents.kind and
    # conversations.kind are int codes (1 = "template"/"default", 2 = "session").
    with raw_engine.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO agents"
                " (workspace_id, id, created_at, name, bundle_location, version, kind)"
                " VALUES (0, 'ag_tmpl', 1, 'my-template', 'ag_tmpl/b', 1, 1),"
                "        (0, 'ag_sess', 2, 'my-session', 'ag_sess/b', 1, 2)"
            )
        )
        conn.execute(
            sa.text(
                "INSERT INTO conversations"
                " (workspace_id, id, created_at, updated_at, root_conversation_id,"
                "  kind, agent_id, title)"
                " VALUES (0, 'conv_1', 3, 3, 'conv_1', 1, 'ag_sess', '')"
            )
        )

    # Downgrade to n1a2b3c4d5e6 (runs o1a2b3c4d5e6 downgrade which restores session_id).
    config2 = _build_alembic_config(uri)
    with raw_engine.begin() as conn:
        config2.attributes["connection"] = conn
        command.downgrade(config2, "n1a2b3c4d5e6")

    # kind must be gone, session_id must be back.
    columns = {c["name"] for c in sa.inspect(raw_engine).get_columns("agents")}
    assert "kind" not in columns
    assert "session_id" in columns

    # The session-scoped agent should have session_id back-populated from
    # conversations.agent_id; the template agent should have NULL.
    with raw_engine.begin() as conn:
        rows = {
            row[0]: row[1]
            for row in conn.execute(sa.text("SELECT id, session_id FROM agents ORDER BY id"))
        }
    assert rows["ag_tmpl"] is None
    assert rows["ag_sess"] == "conv_1"

    raw_engine.dispose()
    clear_engine_cache()