agentic-mobile-control/tests/test_cloud_migrations.py

from __future__ import annotations

import json

import pytest
from sqlalchemy import create_engine, inspect, text

from cloud.schema import (
    HEAD_REVISION,
    SchemaVersionError,
    current_revision,
    downgrade_database,
    require_current_schema,
    upgrade_database,
)


def _database_url(tmp_path) -> str:
    return f"sqlite:///{(tmp_path / 'cloud.sqlite3').as_posix()}"


def test_forward_and_downgrade_migrations_on_empty_database(tmp_path) -> None:
    database_url = _database_url(tmp_path)

    upgrade_database(database_url)

    engine = create_engine(database_url)
    try:
        table_names = set(inspect(engine).get_table_names())
        assert {
            "host_registrations",
            "device_enrollments",
            "pooled_devices",
            "scheduled_tasks",
            "plugins",
            "task_attempts",
            "cloud_users",
            "cloud_user_sessions",
            "cloud_login_throttles",
            "cloud_auth_audit_events",
            "cloud_token_reservations",
            "cloud_token_usage_events",
            "cloud_llm_provider_profiles",
            "cloud_llm_provider_settings",
        } <= table_names
        assert current_revision(database_url) == HEAD_REVISION
        assert (
            connection_scalar(
                engine,
                "select revision from cloud_llm_provider_settings where id = 'global'",
            )
            == 0
        )
        host_columns = {
            column["name"]
            for column in inspect(engine).get_columns("host_registrations")
        }
        assert {
            "agent_instance_id",
            "credential_digest",
            "enrollment_token_digest",
            "revoked_at",
        } <= host_columns
    finally:
        engine.dispose()

    downgrade_database(database_url)

    engine = create_engine(database_url)
    try:
        inspector = inspect(engine)
        assert "device_enrollments" not in inspector.get_table_names()
        assert "task_attempts" not in inspector.get_table_names()
        assert "cloud_users" not in inspector.get_table_names()
        assert "cloud_user_sessions" not in inspector.get_table_names()
        task_columns = {
            column["name"] for column in inspector.get_columns("scheduled_tasks")
        }
        assert "lease_id" not in task_columns
        assert current_revision(database_url) is None
    finally:
        engine.dispose()


def test_legacy_data_survives_upgrade_and_downgrade(tmp_path) -> None:
    database_url = _database_url(tmp_path)
    engine = create_engine(database_url)
    try:
        with engine.begin() as connection:
            _create_legacy_schema(connection)
            connection.execute(
                text(
                    "insert into host_registrations "
                    "(host_id, address, last_seen_at) values "
                    "('host-a', 'local', '2026-01-01T00:00:00+00:00')"
                )
            )
            connection.execute(
                text(
                    "insert into scheduled_tasks "
                    "(id, goal, workflow_definition_id, constraints_json, status, "
                    "assigned_device_id, assigned_host_id, created_at) values "
                    "('task-a', 'goal', null, :constraints, 'queued', null, null, "
                    "'2026-01-01T00:00:00+00:00')"
                ),
                {
                    "constraints": json.dumps(
                        {"driver_type": None, "capability_tags": []}
                    )
                },
            )
    finally:
        engine.dispose()

    upgrade_database(database_url)
    engine = create_engine(database_url)
    try:
        with engine.connect() as connection:
            assert (
                connection.scalar(text("select count(*) from host_registrations")) == 1
            )
            assert connection.scalar(text("select count(*) from scheduled_tasks")) == 1
        task_columns = {
            column["name"] for column in inspect(engine).get_columns("scheduled_tasks")
        }
        assert {"attempt_count", "lease_id", "result_json"} <= task_columns
        host_columns = {
            column["name"]
            for column in inspect(engine).get_columns("host_registrations")
        }
        assert {"agent_instance_id", "credential_digest", "revoked_at"} <= (
            host_columns
        )
        assert "device_enrollments" in inspect(engine).get_table_names()
    finally:
        engine.dispose()

    downgrade_database(database_url)
    engine = create_engine(database_url)
    try:
        with engine.connect() as connection:
            assert (
                connection.scalar(text("select count(*) from host_registrations")) == 1
            )
            assert connection.scalar(text("select count(*) from scheduled_tasks")) == 1
        task_columns = {
            column["name"] for column in inspect(engine).get_columns("scheduled_tasks")
        }
        assert "attempt_count" not in task_columns
    finally:
        engine.dispose()


def test_enrollment_downgrade_to_revision_0001_preserves_legacy_state(
    tmp_path,
) -> None:
    database_url = _database_url(tmp_path)
    upgrade_database(database_url)
    engine = create_engine(database_url)
    try:
        with engine.begin() as connection:
            connection.execute(
                text(
                    "insert into host_registrations "
                    "(host_id, address, last_seen_at) values "
                    "('legacy-host', null, '2026-01-01T00:00:00+00:00')"
                )
            )
    finally:
        engine.dispose()

    downgrade_database(database_url, "0001_cloud_repository")

    engine = create_engine(database_url)
    try:
        inspector = inspect(engine)
        assert "device_enrollments" not in inspector.get_table_names()
        assert connection_scalar(engine, "select count(*) from host_registrations") == 1
        host_columns = {
            column["name"] for column in inspector.get_columns("host_registrations")
        }
        assert "credential_digest" not in host_columns
        assert current_revision(database_url) == "0001_cloud_repository"
    finally:
        engine.dispose()


def test_schema_readiness_requires_head_revision(tmp_path) -> None:
    database_url = _database_url(tmp_path)

    with pytest.raises(SchemaVersionError, match="unversioned"):
        require_current_schema(database_url)

    upgrade_database(database_url)
    require_current_schema(database_url)


def test_planner_decision_log_migration_upgrades_and_downgrades(tmp_path) -> None:
    database_url = _database_url(tmp_path)

    upgrade_database(database_url, "0008_task_progress_columns")
    engine = create_engine(database_url)
    try:
        inspector = inspect(engine)
        assert "planner_decision_log" not in inspector.get_table_names()
    finally:
        engine.dispose()

    upgrade_database(database_url)

    engine = create_engine(database_url)
    try:
        inspector = inspect(engine)
        assert "planner_decision_log" in inspector.get_table_names()
        columns = {col["name"] for col in inspector.get_columns("planner_decision_log")}
        assert {
            "id",
            "host_id",
            "task_id",
            "attempt",
            "step_index",
            "system_prompt",
            "user_prompt",
            "tool_name",
            "arguments_json",
            "created_at",
        } <= columns
        index_names = {
            idx["name"] for idx in inspector.get_indexes("planner_decision_log")
        }
        assert {
            "ix_planner_decision_log_task_attempt",
            "ix_planner_decision_log_host_created",
        } <= index_names
        assert current_revision(database_url) == HEAD_REVISION
    finally:
        engine.dispose()

    downgrade_database(database_url, "0008_task_progress_columns")
    engine = create_engine(database_url)
    try:
        inspector = inspect(engine)
        assert "planner_decision_log" not in inspector.get_table_names()
        assert current_revision(database_url) == "0008_task_progress_columns"
    finally:
        engine.dispose()


def _create_legacy_schema(connection) -> None:
    connection.exec_driver_sql(
        "create table host_registrations ("
        "host_id text primary key, address text, last_seen_at text not null)"
    )
    connection.exec_driver_sql(
        "create table pooled_devices ("
        "device_id text not null, host_id text not null, driver_type text not null, "
        "status text not null, capability_tags_json text not null, synced_at text, "
        "primary key (host_id, device_id))"
    )
    connection.exec_driver_sql(
        "create table scheduled_tasks ("
        "id text primary key, goal text, workflow_definition_id text, "
        "constraints_json text not null, status text not null, "
        "assigned_device_id text, assigned_host_id text, created_at text not null)"
    )
    connection.exec_driver_sql(
        "create table plugins ("
        "name text primary key, version text not null, entry_point_kind text not null, "
        "target text not null, wired integer not null)"
    )


def connection_scalar(engine, statement: str):
    with engine.connect() as connection:
        return connection.scalar(text(statement))