omnigent-ai--omnigent/tests/server/test_identity_migration.py

"""Tests for the accounts → OIDC identity remap.

Covers :func:`omnigent.server.identity_migration.remap_identities` and
``build_domain_mapping`` against a real SQLite database, plus the
``omnigent debug migrate-accounts-to-oidc`` CLI wrapper via Click's
``CliRunner``.

The load-bearing properties: every user-id-bearing column is repointed,
``is_admin`` survives the move, dry runs mutate nothing, grant
collisions merge to the higher level, and an existing distinct NEW id
is refused without ``--force``.
"""

from __future__ import annotations

import pytest
from click.testing import CliRunner
from sqlalchemy import select
from sqlalchemy.orm import Session

from omnigent.cli import cli
from omnigent.db.db_models import (
    SqlAccountToken,
    SqlComment,
    SqlHost,
    SqlPolicy,
)
from omnigent.db.enum_codecs import (
    encode_account_token_kind,
    encode_comment_status,
    encode_host_status,
    encode_policy_scope,
    encode_policy_type,
)
from omnigent.db.utils import get_or_create_engine
from omnigent.server.accounts_store import SqlAlchemyAccountStore
from omnigent.server.identity_migration import build_domain_mapping, remap_identities
from omnigent.server.passwords import hash_password
from omnigent.stores.conversation_store.sqlalchemy_store import SqlAlchemyConversationStore
from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore


def _conversation(db_uri: str) -> str:
    """Create a conversation and return its id (FK target for grants)."""
    return SqlAlchemyConversationStore(db_uri).create_conversation().id


# ── build_domain_mapping ──────────────────────────────────────────


def test_build_domain_mapping_skips_emails_and_reserved(db_uri: str) -> None:
    """Bare usernames map to ``user@domain``; emails / reserved are skipped.

    A user already keyed by email needs no remap, and the ``local`` /
    ``__public__`` sentinels must never be rewritten.
    """
    account_store = SqlAlchemyAccountStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"))
    account_store.create_user_with_password("bob", hash_password("password123"))
    account_store.create_user_with_password("carol@already.com", hash_password("password123"))
    # The "local" sentinel row is created by migrations; ensure it exists.
    SqlAlchemyPermissionStore(db_uri).ensure_user("local", is_admin=True)

    engine = get_or_create_engine(db_uri)
    mapping = build_domain_mapping(engine, "example.com")

    assert mapping == {"alice": "alice@example.com", "bob": "bob@example.com"}


def test_build_domain_mapping_strips_leading_at(db_uri: str) -> None:
    """A ``--domain @example.com`` value is tolerated (leading @ stripped)."""
    SqlAlchemyAccountStore(db_uri).create_user_with_password("alice", hash_password("password123"))
    mapping = build_domain_mapping(get_or_create_engine(db_uri), "@example.com")
    assert mapping == {"alice": "alice@example.com"}


# ── remap_identities: the core move ───────────────────────────────


def test_remap_moves_user_grant_and_admin(db_uri: str) -> None:
    """A committed remap moves the user row + grant and preserves is_admin."""
    account_store = SqlAlchemyAccountStore(db_uri)
    perm_store = SqlAlchemyPermissionStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"), is_admin=True)
    conv_id = _conversation(db_uri)
    perm_store.grant("alice", conv_id, level=3)

    report = remap_identities(
        get_or_create_engine(db_uri),
        {"alice": "alice@example.com"},
        dry_run=False,
    )

    assert report.committed is True
    # Old principal is gone; new one exists and kept admin.
    assert account_store.get_user("alice") is None
    new_user = account_store.get_user("alice@example.com")
    assert new_user is not None
    assert new_user.is_admin is True
    # The grant moved to the new id at the same level.
    moved = perm_store.get("alice@example.com", conv_id)
    assert moved is not None and moved.level == 3
    assert perm_store.get("alice", conv_id) is None


def test_remap_repoints_comments_policies_tokens_hosts(db_uri: str) -> None:
    """Every user-id-bearing column is repointed, not just users/grants."""
    account_store = SqlAlchemyAccountStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"))
    engine = get_or_create_engine(db_uri)

    # Seed one row per referencing table keyed on "alice".
    with Session(engine) as s:
        s.add(
            SqlComment(
                id="cmt_1",
                conversation_id="conv_x",
                path="a.py",
                start_index=0,
                end_index=1,
                body="hi",
                status=encode_comment_status("draft"),
                created_at=1,
                # created_at scaled to epoch-µs, matching the store's invariant.
                updated_at=1_000_000,
                created_by="alice",
            )
        )
        s.add(
            SqlPolicy(
                id="pol_1",
                name="p",
                session_id=None,
                scope=encode_policy_scope("default"),
                created_at=1,
                type=encode_policy_type("python"),
                handler="x.y",
                created_by="alice",
            )
        )
        s.add(
            SqlAccountToken(
                id="tok_1",
                kind=encode_account_token_kind("invite"),
                user_id=None,
                created_by="alice",
                created_at=1,
                expires_at=10_000_000_000,
            )
        )
        s.add(
            SqlHost(
                owner="alice",
                name="laptop",
                host_id="h1",
                status=encode_host_status("offline"),
                created_at=1,
                updated_at=1,
            )
        )
        s.commit()

    remap_identities(engine, {"alice": "alice@example.com"}, dry_run=False)

    with Session(engine) as s:
        assert s.get(SqlComment, (0, "cmt_1")).created_by == "alice@example.com"
        assert s.get(SqlPolicy, (0, "pol_1")).created_by == "alice@example.com"
        assert s.get(SqlAccountToken, (0, "tok_1")).created_by == "alice@example.com"
        host_owners = s.execute(select(SqlHost.owner)).scalars().all()
        assert host_owners == ["alice@example.com"]


def test_dry_run_mutates_nothing_but_reports(db_uri: str) -> None:
    """A dry run reports would-change counts but leaves the DB untouched."""
    account_store = SqlAlchemyAccountStore(db_uri)
    perm_store = SqlAlchemyPermissionStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"), is_admin=True)
    conv_id = _conversation(db_uri)
    perm_store.grant("alice", conv_id, level=3)

    report = remap_identities(
        get_or_create_engine(db_uri),
        {"alice": "alice@example.com"},
        dry_run=True,
    )

    assert report.committed is False
    # Report still reflects what *would* change.
    assert report.per_table.get("session_permissions") == 1
    assert report.per_table.get("users") == 1
    # But nothing actually moved.
    assert account_store.get_user("alice") is not None
    assert account_store.get_user("alice@example.com") is None
    assert perm_store.get("alice", conv_id) is not None


def test_grant_collision_merges_to_higher_level(db_uri: str) -> None:
    """When NEW already has a grant on the same conversation, levels merge to max."""
    account_store = SqlAlchemyAccountStore(db_uri)
    perm_store = SqlAlchemyPermissionStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"))
    account_store.create_user_with_password("alice@example.com", hash_password("password123"))
    conv_id = _conversation(db_uri)
    perm_store.grant("alice", conv_id, level=3)  # old has manage
    perm_store.grant("alice@example.com", conv_id, level=1)  # new has read

    # NEW exists distinctly → needs force to merge.
    report = remap_identities(
        get_or_create_engine(db_uri),
        {"alice": "alice@example.com"},
        dry_run=False,
        force=True,
    )

    assert report.committed is True
    merged = perm_store.get("alice@example.com", conv_id)
    assert merged is not None and merged.level == 3  # max(3, 1)
    assert perm_store.get("alice", conv_id) is None


def test_refuses_existing_new_without_force(db_uri: str) -> None:
    """Mapping onto an existing distinct NEW id is refused unless --force."""
    account_store = SqlAlchemyAccountStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"))
    account_store.create_user_with_password("alice@example.com", hash_password("password123"))

    report = remap_identities(
        get_or_create_engine(db_uri),
        {"alice": "alice@example.com"},
        dry_run=False,
        force=False,
    )

    assert report.refused == ["alice -> alice@example.com"]
    # Both rows untouched.
    assert account_store.get_user("alice") is not None
    assert account_store.get_user("alice@example.com") is not None


def test_skipped_missing_old_user(db_uri: str) -> None:
    """An old id with no users row is recorded in skipped_missing."""
    report = remap_identities(
        get_or_create_engine(db_uri),
        {"ghost": "ghost@example.com"},
        dry_run=False,
    )
    assert report.skipped_missing == ["ghost"]
    assert report.per_table == {}


# ── CLI wrapper ───────────────────────────────────────────────────


def test_cli_dry_run_by_default(db_uri: str) -> None:
    """``migrate-to-oidc`` without --commit is a dry run that changes nothing."""
    account_store = SqlAlchemyAccountStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"), is_admin=True)

    result = CliRunner().invoke(
        cli,
        ["debug", "migrate-accounts-to-oidc", db_uri, "--domain", "example.com"],
    )

    assert result.exit_code == 0, result.output
    assert "DRY RUN" in result.output
    assert "alice  ->  alice@example.com" in result.output
    # Surfaces the IdP-email-mismatch reminder so the operator verifies
    # the targets match what their IdP returns before committing.
    assert "must match the email your IdP returns" in result.output
    # No --commit → unchanged.
    assert account_store.get_user("alice") is not None
    assert account_store.get_user("alice@example.com") is None


def test_cli_commit_applies(db_uri: str) -> None:
    """``--commit`` applies the remap."""
    account_store = SqlAlchemyAccountStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"), is_admin=True)

    result = CliRunner().invoke(
        cli,
        ["debug", "migrate-accounts-to-oidc", db_uri, "--domain", "example.com", "--commit"],
    )

    assert result.exit_code == 0, result.output
    assert "COMMITTED" in result.output
    assert account_store.get_user("alice") is None
    assert account_store.get_user("alice@example.com") is not None


def test_cli_requires_a_mapping(db_uri: str) -> None:
    """With neither --domain nor --map, the command errors (nothing to do)."""
    result = CliRunner().invoke(cli, ["debug", "migrate-accounts-to-oidc", db_uri])
    assert result.exit_code != 0
    assert "nothing to migrate" in result.output


def test_cli_map_overrides_domain(db_uri: str) -> None:
    """An explicit --map pair wins over the --domain-derived mapping."""
    account_store = SqlAlchemyAccountStore(db_uri)
    account_store.create_user_with_password("alice", hash_password("password123"))

    result = CliRunner().invoke(
        cli,
        [
            "debug",
            "migrate-accounts-to-oidc",
            db_uri,
            "--domain",
            "example.com",
            "--map",
            "alice=alice@corp.com",
            "--commit",
        ],
    )

    assert result.exit_code == 0, result.output
    assert account_store.get_user("alice@corp.com") is not None
    assert account_store.get_user("alice@example.com") is None


@pytest.mark.parametrize("bad", ["aliceonly", "=new", "old="])
def test_cli_rejects_malformed_map(db_uri: str, bad: str) -> None:
    """``--map`` without a valid OLD=NEW shape is rejected."""
    result = CliRunner().invoke(cli, ["debug", "migrate-accounts-to-oidc", db_uri, "--map", bad])
    assert result.exit_code != 0