cleveragents-core/robot/helper_resource_registry_migration.py

"""Helper utilities for resource registry migration Robot smoke tests.

Validates that ``Base.metadata.create_all`` (the schema-creation path used
by ``init_database``) produces the three resource-registry tables:

* ``resource_types``
* ``resources``
* ``resource_edges``

Each command prints a single ``<command>-ok`` token on success so the
calling Robot test can assert on ``stdout``.
"""

from __future__ import annotations

import sys
import tempfile
from pathlib import Path

from sqlalchemy import inspect

from cleveragents.infrastructure.database.models import init_database


def _schema_creates_tables() -> None:
    """Verify that init_database creates resource registry tables."""
    tmp = tempfile.mktemp(suffix=".db")
    db_url = f"sqlite:///{tmp}"
    try:
        engine = init_database(db_url)
        inspector = inspect(engine)
        tables = inspector.get_table_names()
        expected = ["resource_types", "resources", "resource_edges"]
        for table in expected:
            assert table in tables, f"Table {table} not found. Tables: {tables}"
        engine.dispose()
        print("schema-creates-tables-ok")
    finally:
        Path(tmp).unlink(missing_ok=True)


def _table_has_expected_columns() -> None:
    """Verify core columns exist on each resource registry table."""
    tmp = tempfile.mktemp(suffix=".db")
    db_url = f"sqlite:///{tmp}"
    try:
        engine = init_database(db_url)
        inspector = inspect(engine)

        # resource_types columns
        rt_cols = {c["name"] for c in inspector.get_columns("resource_types")}
        for col in (
            "name",
            "namespace",
            "resource_kind",
            "user_addable",
            "created_at",
            "updated_at",
        ):
            assert col in rt_cols, f"Column {col} missing from resource_types"

        # resources columns
        r_cols = {c["name"] for c in inspector.get_columns("resources")}
        for col in (
            "resource_id",
            "type_name",
            "resource_kind",
            "created_at",
            "updated_at",
        ):
            assert col in r_cols, f"Column {col} missing from resources"

        # resource_edges columns
        re_cols = {c["name"] for c in inspector.get_columns("resource_edges")}
        for col in ("parent_id", "child_id", "link_type", "created_at"):
            assert col in re_cols, f"Column {col} missing from resource_edges"

        engine.dispose()
        print("table-columns-ok")
    finally:
        Path(tmp).unlink(missing_ok=True)


def _migration_idempotent() -> None:
    """Verify that calling init_database twice is safe (idempotent)."""
    tmp = tempfile.mktemp(suffix=".db")
    db_url = f"sqlite:///{tmp}"
    try:
        engine1 = init_database(db_url)
        engine1.dispose()

        # Second call should not raise
        engine2 = init_database(db_url)
        inspector = inspect(engine2)
        tables = inspector.get_table_names()
        assert "resource_types" in tables
        assert "resources" in tables
        assert "resource_edges" in tables
        engine2.dispose()
        print("migration-idempotent-ok")
    finally:
        Path(tmp).unlink(missing_ok=True)


def main() -> None:
    """Dispatch to the requested sub-command."""
    if len(sys.argv) < 2:
        raise SystemExit("Expected command argument")
    command = sys.argv[1]
    commands = {
        "schema-creates-tables": _schema_creates_tables,
        "table-columns": _table_has_expected_columns,
        "migration-idempotent": _migration_idempotent,
    }
    if command not in commands:
        raise SystemExit(f"Unknown command: {command}")
    commands[command]()


if __name__ == "__main__":
    main()