DataOps-platform/tests/integration/test_phase3_wp02_enterprise_identity_postgres.py

from __future__ import annotations

import os
import threading
import uuid
from datetime import UTC, datetime, timedelta

import pytest
from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session

from app.core.system.enterprise_identity import (
    DirectorySynchronizer,
    EmergencyAccess,
    IdentityPolicyError,
    sanitize_audit_detail,
)
from app.core.system.identity_repository import PostgresIdentityRepository
from app.core.system.identity_sessions import SessionManager

pytestmark = pytest.mark.integration


def _uid() -> str:
    return str(uuid.uuid4())


def _identity_attributes(username: str, roles: list[str] | None = None) -> dict:
    return {
        "provider_uid": "filled-by-synchronizer",
        "username": username,
        "display_name": username.title(),
        "department": "engineering",
        "groups": ["data-stewards"],
        "roles": roles or ["viewer"],
        "mapping_version": "integration-v1",
        "claims_digest": "a" * 64,
        "authorization_scope": {"environments": ["test"]},
    }


@pytest.fixture
def identity_database():
    database_url = os.environ.get("TEST_DATABASE_URL")
    if not database_url:
        pytest.skip("TEST_DATABASE_URL is required")
    engine = create_engine(database_url, pool_pre_ping=True)
    marker = uuid.uuid4().hex[:10]
    users = {name: _uid() for name in ("requester", "breakglass", "approver_a", "approver_b")}
    providers: list[str] = []
    with engine.begin() as connection:
        for name, user_uid in users.items():
            connection.execute(text("""
                INSERT INTO public.users(id,username,display_name,password_hash,status)
                VALUES(CAST(:uid AS uuid),:username,:display,'integration-only','active')
            """), {"uid": user_uid, "username": f"p3wp02-{name}-{marker}", "display": name})
            connection.execute(text("""
                INSERT INTO public.user_roles(user_id,role_id)
                SELECT CAST(:uid AS uuid),id FROM public.roles WHERE name='admin'
            """), {"uid": user_uid})
    state = {"engine": engine, "users": users, "providers": providers, "marker": marker}
    try:
        yield state
    finally:
        with engine.begin() as connection:
            linked_users: set[str] = set()
            for provider in providers:
                linked_users.update(str(value) for value in connection.execute(text(
                    "SELECT user_uid FROM public.enterprise_identity_links WHERE provider_uid=CAST(:provider AS uuid)"
                ), {"provider": provider}).scalars())
                connection.execute(text("DELETE FROM public.identity_audit_events WHERE provider_uid=CAST(:provider AS uuid)"),
                                   {"provider": provider})
                connection.execute(text("DELETE FROM public.identity_directory_conflicts WHERE provider_uid=CAST(:provider AS uuid)"),
                                   {"provider": provider})
                connection.execute(text("DELETE FROM public.identity_directory_events WHERE provider_uid=CAST(:provider AS uuid)"),
                                   {"provider": provider})
                connection.execute(text("DELETE FROM public.identity_directory_checkpoints WHERE provider_uid=CAST(:provider AS uuid)"),
                                   {"provider": provider})
                connection.execute(text("DELETE FROM public.identity_organization_nodes WHERE provider_uid=CAST(:provider AS uuid)"),
                                   {"provider": provider})
            all_users = set(users.values()) | linked_users
            for user_uid in all_users:
                connection.execute(text("DELETE FROM public.identity_exchange_codes WHERE session_uid IN (SELECT uid FROM public.identity_sessions WHERE user_uid=CAST(:uid AS uuid))"),
                                   {"uid": user_uid})
                connection.execute(text("DELETE FROM public.identity_sessions WHERE user_uid=CAST(:uid AS uuid)"), {"uid": user_uid})
            for request_uid in connection.execute(text("SELECT uid::text FROM public.identity_emergency_requests WHERE requester_uid=CAST(:uid AS uuid)"),
                                                  {"uid": users["requester"]}).scalars().all():
                connection.execute(text("DELETE FROM public.identity_emergency_approvals WHERE request_uid=CAST(:uid AS uuid)"), {"uid": request_uid})
                connection.execute(text("DELETE FROM public.identity_emergency_requests WHERE uid=CAST(:uid AS uuid)"), {"uid": request_uid})
            for provider in providers:
                connection.execute(text("DELETE FROM public.enterprise_identity_links WHERE provider_uid=CAST(:provider AS uuid)"),
                                   {"provider": provider})
            for user_uid in all_users:
                connection.execute(text("DELETE FROM public.identity_audit_events WHERE actor_uid=CAST(:uid AS uuid)"), {"uid": user_uid})
                connection.execute(text("DELETE FROM public.user_roles WHERE user_id=CAST(:uid AS uuid)"), {"uid": user_uid})
                connection.execute(text("DELETE FROM public.users WHERE id=CAST(:uid AS uuid)"), {"uid": user_uid})
            assert connection.execute(
                text("SELECT count(*) FROM public.users WHERE username LIKE :pattern"),
                {"pattern": f"%-{marker}"},
            ).scalar_one() == 0
        engine.dispose()


def _provider(state) -> str:
    provider = _uid()
    state["providers"].append(provider)
    return provider


def test_real_470_schema_constraints_and_provider_scoped_directory_unit_of_work(identity_database, monkeypatch):
    engine = identity_database["engine"]
    provider_a, provider_b = _provider(identity_database), _provider(identity_database)
    with Session(engine) as session:
        assert session.execute(text("SELECT version_num FROM public.alembic_version")).scalar_one() == "20260802_471"
        tables = set(session.execute(text("""
            SELECT table_name FROM information_schema.tables WHERE table_schema='public'
              AND table_name LIKE 'identity_%'
        """)).scalars())
        assert {"identity_directory_events", "identity_directory_checkpoints", "identity_directory_conflicts",
                "identity_organization_nodes", "identity_sessions", "identity_emergency_requests"} <= tables
        constraints = " ".join(session.execute(text("""
            SELECT pg_get_constraintdef(oid) FROM pg_constraint
            WHERE conrelid IN ('public.identity_sessions'::regclass,
                                'public.identity_directory_events'::regclass,
                                'public.identity_emergency_requests'::regclass,
                                'public.identity_idp_config_versions'::regclass)
        """)).scalars())
        assert "UNIQUE (refresh_hash)" in constraints
        assert "UNIQUE (provider_uid, source, source_event_id)" in constraints
        assert "expires_at > starts_at" in constraints
        assert "DATAOPS_OIDC_" in constraints

        repo = PostgresIdentityRepository(session)
        sync = DirectorySynchronizer(repo, SessionManager(repo, secret="integration-secret"))
        first = sync.apply(provider_uid=provider_a, source="scim", source_event_id="join-a", cursor="1",
                           cursor_sequence=1, event_type="JOINER", subject="shared-subject",
                           attributes=_identity_attributes(f"alice-a-{identity_database['marker']}"))
        second = sync.apply(provider_uid=provider_b, source="scim", source_event_id="join-b", cursor="1",
                            cursor_sequence=1, event_type="JOINER", subject="shared-subject",
                            attributes=_identity_attributes(f"alice-b-{identity_database['marker']}"))
        assert first["user_uid"] != second["user_uid"]
        assert repo.get_identity(provider_a, "shared-subject")["user_uid"] == first["user_uid"]
        assert repo.get_identity(provider_b, "shared-subject")["user_uid"] == second["user_uid"]

        session_a = SessionManager(repo, secret="integration-secret").create(
            provider_uid=provider_a, subject="shared-subject", user_uid=first["user_uid"],
            roles=["viewer"], identity_source="oidc")
        session_b = SessionManager(repo, secret="integration-secret").create(
            provider_uid=provider_b, subject="shared-subject", user_uid=second["user_uid"],
            roles=["viewer"], identity_source="oidc")
        assert {row["uid"] for row in repo.list_sessions(provider_a, "shared-subject")} == {session_a.session_uid}
        assert {row["uid"] for row in repo.list_sessions(provider_b, "shared-subject")} == {session_b.session_uid}

        sync.apply(provider_uid=provider_a, source="scim", source_event_id="move-a", cursor="2",
                   cursor_sequence=2, event_type="MOVER", subject="shared-subject",
                   attributes=_identity_attributes(f"alice-a-{identity_database['marker']}", ["editor"]))
        assert repo.get_session(session_a.session_uid)["status"] == "revoked"
        assert repo.get_session(session_b.session_uid)["status"] == "active"
        assert repo.get_directory_checkpoint(provider_a, "scim") == {"cursor": "2", "cursor_sequence": 2}
        with pytest.raises(IdentityPolicyError, match="idempotency conflict"):
            sync.apply(provider_uid=provider_a, source="scim", source_event_id="move-a", cursor="2",
                       cursor_sequence=2, event_type="MOVER", subject="shared-subject",
                       attributes=_identity_attributes(f"changed-{identity_database['marker']}", ["editor"]))
        assert session.execute(text("SELECT count(*) FROM public.identity_directory_conflicts WHERE provider_uid=CAST(:p AS uuid)"),
                               {"p": provider_a}).scalar_one() == 1

        rollback_subject = "rollback-subject"
        original_insert = repo._insert_directory_event

        def fail_after_mutation(event, result):
            raise RuntimeError("injected event persistence failure")

        monkeypatch.setattr(repo, "_insert_directory_event", fail_after_mutation)
        with pytest.raises(RuntimeError, match="injected"):
            sync.apply(provider_uid=provider_a, source="rollback", source_event_id="rollback-1", cursor="1",
                       cursor_sequence=1, event_type="JOINER", subject=rollback_subject,
                       attributes=_identity_attributes(f"rollback-{identity_database['marker']}"))
        monkeypatch.setattr(repo, "_insert_directory_event", original_insert)
        assert repo.get_identity(provider_a, rollback_subject) is None
        assert repo.get_directory_checkpoint(provider_a, "rollback") is None


def test_organization_node_incremental_lifecycle_is_provider_isolated_and_auditable(identity_database):
    engine = identity_database["engine"]
    provider_a, provider_b = _provider(identity_database), _provider(identity_database)
    with Session(engine) as session:
        repo = PostgresIdentityRepository(session)
        sync = DirectorySynchronizer(repo, SessionManager(repo, secret="integration-secret"))
        department = {"external_id": "dept-ops", "display_name": "Operations", "action": "UPSERT",
                      "attributes": {"cost_center": "CC-1"}}
        created = sync.apply(provider_uid=provider_a, source="directory", source_event_id="dept-1", cursor="1",
                             cursor_sequence=1, event_type="DEPARTMENT", attributes=department)
        assert created["status"] == "active"
        disabled = sync.apply(provider_uid=provider_a, source="directory", source_event_id="dept-2", cursor="2",
                              cursor_sequence=2, event_type="DEPARTMENT",
                              attributes={"external_id": "dept-ops", "action": "DISABLE"})
        assert disabled["status"] == "disabled"
        restored = sync.apply(provider_uid=provider_a, source="directory", source_event_id="dept-3", cursor="3",
                              cursor_sequence=3, event_type="DEPARTMENT",
                              attributes={**department, "action": "RESTORE", "display_name": "Operations Restored"})
        assert restored["status"] == "active"
        assert sync.apply(provider_uid=provider_a, source="directory", source_event_id="dept-3", cursor="3",
                          cursor_sequence=3, event_type="DEPARTMENT",
                          attributes={**department, "action": "RESTORE", "display_name": "Operations Restored"})["idempotent"] is True
        sync.apply(provider_uid=provider_a, source="directory", source_event_id="group-1", cursor="4",
                   cursor_sequence=4, event_type="GROUP",
                   attributes={"external_id": "group-stewards", "display_name": "Stewards", "action": "UPSERT"})
        sync.apply(provider_uid=provider_b, source="directory", source_event_id="dept-other", cursor="1",
                   cursor_sequence=1, event_type="DEPARTMENT", attributes=department)
        node_a = repo.get_organization_node(provider_a, "department", "dept-ops")
        node_b = repo.get_organization_node(provider_b, "department", "dept-ops")
        assert node_a["display_name"] == "Operations Restored"
        assert node_b["display_name"] == "Operations"
        repo.audit({"event_type": "directory_delta", "outcome": "success", "provider_uid": provider_a,
                    "resource_type": "organization_node", "resource_uid": created["node_uid"],
                    "safe_detail": sanitize_audit_detail({"event_type": "DEPARTMENT", "token": "must-not-persist"})})
        audit = session.execute(text("SELECT safe_detail FROM public.identity_audit_events WHERE provider_uid=CAST(:p AS uuid)"),
                                {"p": provider_a}).scalar_one()
        assert audit == {"event_type": "DEPARTMENT"}


def test_refresh_rotation_and_session_limit_are_atomic_under_concurrency(identity_database):
    engine = identity_database["engine"]
    provider = _provider(identity_database)
    with Session(engine) as session:
        repo = PostgresIdentityRepository(session)
        sync = DirectorySynchronizer(repo, SessionManager(repo, secret="integration-secret"))
        identity = sync.apply(provider_uid=provider, source="scim", source_event_id="join", cursor="1",
                              cursor_sequence=1, event_type="JOINER", subject="concurrent-subject",
                              attributes=_identity_attributes(f"concurrent-{identity_database['marker']}"))
        original = SessionManager(repo, secret="integration-secret").create(
            provider_uid=provider, subject="concurrent-subject", user_uid=identity["user_uid"],
            roles=["viewer"], identity_source="oidc")
        family_uid = repo.get_session(original.session_uid)["family_uid"]

    barrier = threading.Barrier(2)
    outcomes: list[str] = []
    outcome_lock = threading.Lock()

    def rotate_once():
        with Session(engine) as thread_session:
            manager = SessionManager(PostgresIdentityRepository(thread_session), secret="integration-secret")
            barrier.wait()
            try:
                manager.refresh(original.refresh_token)
                outcome = "rotated"
            except IdentityPolicyError as exc:
                outcome = str(exc)
            with outcome_lock:
                outcomes.append(outcome)

    threads = [threading.Thread(target=rotate_once) for _ in range(2)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=15)
    assert all(not thread.is_alive() for thread in threads)
    assert outcomes.count("rotated") <= 1
    assert "refresh reuse detected" in outcomes
    with Session(engine) as session:
        statuses = session.execute(text("SELECT status FROM public.identity_sessions WHERE family_uid=CAST(:family AS uuid)"),
                                   {"family": family_uid}).scalars().all()
        assert statuses and set(statuses) == {"revoked"}

    create_barrier = threading.Barrier(6)
    create_errors: list[Exception] = []

    def create_once(index: int):
        with Session(engine) as thread_session:
            manager = SessionManager(PostgresIdentityRepository(thread_session), secret="integration-secret", max_sessions=2)
            create_barrier.wait()
            try:
                manager.create(provider_uid=provider, subject="concurrent-subject", user_uid=identity["user_uid"],
                               roles=["viewer"], identity_source="oidc")
            except Exception as exc:  # captured and asserted below
                create_errors.append(exc)

    creators = [threading.Thread(target=create_once, args=(index,)) for index in range(6)]
    for thread in creators:
        thread.start()
    for thread in creators:
        thread.join(timeout=15)
    assert not create_errors
    with Session(engine) as session:
        active_count = session.execute(text("SELECT count(*) FROM public.identity_sessions WHERE user_uid=CAST(:u AS uuid) AND status='active'"),
                                       {"u": identity["user_uid"]}).scalar_one()
        assert active_count <= 2


def test_emergency_session_rechecks_local_admin_and_close_revokes(identity_database):
    engine = identity_database["engine"]
    users = identity_database["users"]
    now = datetime.now(UTC)
    with Session(engine) as session:
        repo = PostgresIdentityRepository(session)
        service = EmergencyAccess(repo, clock=lambda: now)
        request = service.request(requester_uid=users["requester"], account_uid=users["breakglass"],
                                  reason="integration outage", expires_at=now + timedelta(minutes=30),
                                  account_is_local_active_admin=True)
        service.approve(request["uid"], approver_uid=users["approver_a"])
        service.approve(request["uid"], approver_uid=users["approver_b"])
        active = service.activate(request["uid"])
        manager = SessionManager(repo, secret="integration-secret", clock=lambda: now)
        credentials = manager.create(subject=users["breakglass"], user_uid=users["breakglass"], roles=["admin"],
                                     identity_source="emergency", emergency_request_uid=request["uid"],
                                     emergency_expires_at=active["expires_at"])
        session.execute(text("DELETE FROM public.user_roles WHERE user_id=CAST(:u AS uuid)"), {"u": users["breakglass"]})
        session.commit()
        with pytest.raises(IdentityPolicyError, match="active local administrator"):
            manager.create(subject=users["breakglass"], user_uid=users["breakglass"], roles=["admin"],
                           identity_source="emergency", emergency_request_uid=request["uid"],
                           emergency_expires_at=active["expires_at"])
        session.execute(text("INSERT INTO public.user_roles(user_id,role_id) SELECT CAST(:u AS uuid),id FROM public.roles WHERE name='admin'"),
                        {"u": users["breakglass"]})
        session.commit()
        service.close(request["uid"], actor_uid=users["approver_a"])
        assert repo.get_session(credentials.session_uid)["status"] == "revoked"
        assert repo.get_session(credentials.session_uid)["emergency_request_uid"] == request["uid"]


def test_oidc_login_persistence_rolls_back_identity_session_and_exchange_together(identity_database, monkeypatch):
    engine = identity_database["engine"]
    provider = _provider(identity_database)
    user_uid = _uid()
    subject = "atomic-callback-subject"
    with Session(engine) as session:
        repo = PostgresIdentityRepository(session)
        repo.put_identity(provider, subject, {
            "provider_uid": provider, "user_uid": user_uid,
            "username": f"atomic-{identity_database['marker']}", "display_name": "Atomic Callback",
            "department": "security", "groups": ["stewards"], "roles": ["viewer"],
            "authorization_scope": {"environments": ["test"]}, "mapping_version": "integration-v1",
            "claims_digest": "b" * 64, "status": "active", "token_version": 1,
        }, commit=False)
        credentials = SessionManager(repo, secret="integration-secret").create(
            provider_uid=provider, subject=subject, user_uid=user_uid, roles=["viewer"],
            identity_source="oidc", commit=False)

        def fail_exchange(record, *, commit=True):
            del record, commit
            raise RuntimeError("injected exchange persistence failure")

        monkeypatch.setattr(repo, "insert_exchange_code", fail_exchange)
        with pytest.raises(RuntimeError, match="injected"):
            repo.insert_exchange_code({"uid": _uid(), "code_hash": "c" * 64,
                                       "session_uid": credentials.session_uid,
                                       "redirect_uri": "https://dataops.example.com/login/callback",
                                       "expires_at": datetime.now(UTC) + timedelta(minutes=2)}, commit=False)
        session.rollback()
        assert repo.get_identity(provider, subject) is None
        assert repo.get_session(credentials.session_uid) is None
        assert session.execute(text("SELECT count(*) FROM public.users WHERE id=CAST(:u AS uuid)"),
                               {"u": user_uid}).scalar_one() == 0