stack/tests/aco/test_lake_governance.py

"""Tests for aco.lake.governance — declarative grant management."""

from __future__ import annotations

from unittest.mock import MagicMock

from aco.lake.governance import (
    GovernancePolicy,
    GrantRule,
    _resolve_schemas,
    apply_governance,
    audit_governance,
)


class TestResolveSchemas:
    def test_wildcard_expands(self):
        rule = GrantRule(group="g", schemas=["*"], privileges=["SELECT"])
        assert _resolve_schemas(rule, ["core", "claims"]) == ["core", "claims"]

    def test_explicit_schemas(self):
        rule = GrantRule(group="g", schemas=["core"], privileges=["SELECT"])
        assert _resolve_schemas(rule, ["core", "claims"]) == ["core"]

    def test_missing_schema_filtered(self):
        rule = GrantRule(group="g", schemas=["missing"], privileges=["SELECT"])
        assert _resolve_schemas(rule, ["core"]) == []


class TestGovernancePolicyModel:
    def test_basic_construction(self):
        policy = GovernancePolicy(
            catalog="aco",
            rules=[
                GrantRule(
                    group="analysts",
                    schemas=["readmissions"],
                    privileges=["SELECT"],
                ),
            ],
            phi_schemas=["core"],
        )
        assert policy.catalog == "aco"
        assert len(policy.rules) == 1
        assert policy.phi_schemas == ["core"]

    def test_empty_rules(self):
        policy = GovernancePolicy(catalog="aco")
        assert policy.rules == []


class TestApplyGovernance:
    def test_dry_run_reports_actions(self):
        client = MagicMock()
        # Mock schemas list
        schema_mock = MagicMock()
        schema_mock.name = "readmissions"
        client._ws.schemas.list.return_value = [schema_mock]

        policy = GovernancePolicy(
            catalog="aco",
            rules=[
                GrantRule(
                    group="analysts",
                    schemas=["readmissions"],
                    privileges=["SELECT"],
                ),
            ],
        )

        actions = apply_governance(client, policy, dry_run=True)
        assert len(actions) == 1
        assert "[dry-run]" in actions[0]
        assert "GRANT SELECT" in actions[0]
        assert "analysts" in actions[0]
        # Should NOT call ws.grants.update in dry_run
        client._ws.grants.update.assert_not_called()

    def test_apply_calls_sdk(self):
        client = MagicMock()
        schema_mock = MagicMock()
        schema_mock.name = "core"
        client._ws.schemas.list.return_value = [schema_mock]

        policy = GovernancePolicy(
            catalog="aco",
            rules=[
                GrantRule(
                    group="pipeline_svc",
                    schemas=["core"],
                    privileges=["ALL_PRIVILEGES"],
                ),
            ],
        )

        actions = apply_governance(client, policy, dry_run=False)
        assert len(actions) == 1
        assert "GRANT" in actions[0]
        client._ws.grants.update.assert_called_once()

    def test_wildcard_expands_to_all(self):
        client = MagicMock()
        s1, s2 = MagicMock(), MagicMock()
        s1.name, s2.name = "core", "claims"
        client._ws.schemas.list.return_value = [s1, s2]

        policy = GovernancePolicy(
            catalog="aco",
            rules=[
                GrantRule(group="g", schemas=["*"], privileges=["SELECT"]),
            ],
        )

        actions = apply_governance(client, policy, dry_run=True)
        assert len(actions) == 2


class TestAuditGovernance:
    def test_detects_missing_grants(self):
        client = MagicMock()
        schema_mock = MagicMock()
        schema_mock.name = "core"
        client._ws.schemas.list.return_value = [schema_mock]

        # No grants exist
        resp = MagicMock()
        resp.privilege_assignments = []
        client._ws.grants.get.return_value = resp

        policy = GovernancePolicy(
            catalog="aco",
            rules=[
                GrantRule(group="analysts", schemas=["core"], privileges=["SELECT"]),
            ],
        )

        drifts = audit_governance(client, policy)
        assert len(drifts) == 1
        assert drifts[0].principal == "analysts"
        assert drifts[0].action == "grant"

    def test_no_drift_when_matching(self):
        client = MagicMock()
        schema_mock = MagicMock()
        schema_mock.name = "core"
        client._ws.schemas.list.return_value = [schema_mock]

        # Grants match exactly
        priv = MagicMock()
        priv.privilege = MagicMock()
        priv.privilege.value = "SELECT"
        assignment = MagicMock()
        assignment.principal = "analysts"
        assignment.privileges = [priv]
        resp = MagicMock()
        resp.privilege_assignments = [assignment]
        client._ws.grants.get.return_value = resp

        policy = GovernancePolicy(
            catalog="aco",
            rules=[
                GrantRule(group="analysts", schemas=["core"], privileges=["SELECT"]),
            ],
        )

        drifts = audit_governance(client, policy)
        assert len(drifts) == 0