agentic-mobile-control/tests/test_cloud_governance.py

from __future__ import annotations

from datetime import UTC, datetime

import pytest

from cloud.auth import Principal
from cloud.config import CloudConfig
from cloud.database import CloudDatabase
from cloud.governance import enforce_user_submission_policy
from cloud.plugins import PluginRegistry
from cloud.pool import DevicePool
from cloud.sdk.api import create_cloud_router
from cloud.sdk.governance_api import create_governance_router
from cloud.scheduler import TaskScheduler
from cloud.user_auth import UserAccount
from core.models import Device


pytest.importorskip("fastapi")
from fastapi import FastAPI  # noqa: E402
from fastapi.testclient import TestClient  # noqa: E402


class _PrincipalProvider:
    def __init__(self, principal: Principal) -> None:
        self.principal = principal

    def authenticate(self, _request: object) -> Principal:
        return self.principal


def _config() -> CloudConfig:
    return CloudConfig(max_queue_depth=20)


def _add_user(repository, user_id: str = "user-a") -> None:
    now = datetime.now(UTC)
    repository.create_user(
        UserAccount(
            id=user_id,
            username=user_id,
            username_normalized=user_id,
            display_name=user_id,
            role="operator",
            enabled=True,
            must_change_password=False,
            authentication_version=1,
            created_at=now,
            updated_at=now,
            password_hash="not-used-by-this-test",
        )
    )


def test_user_submission_policy_requires_permitted_explicit_target(tmp_path) -> None:
    database = CloudDatabase(f"sqlite:///{(tmp_path / 'governance.sqlite3').as_posix()}")
    try:
        _add_user(database.repository)
        database.repository.upsert_user_submission_policy(
            user_id="user-a",
            submission_enabled=True,
            allowed_host_ids=("host-a",),
            allowed_device_targets=(("host-a", "device-a"),),
            updated_at=datetime.now(UTC),
        )
        pool = DevicePool(database.repository, _config())
        pool.sync_host_devices(
            "host-a",
            [Device(id="device-a", driver_type="wda", status="idle")],  # type: ignore[arg-type]
        )
        scheduler = TaskScheduler(pool, database.repository, _config())
        app = FastAPI()
        app.include_router(
            create_cloud_router(
                pool=pool,
                scheduler=scheduler,
                plugin_registry=PluginRegistry(database.repository),
                auth_provider=_PrincipalProvider(
                    Principal(id="user:user-a", scopes=frozenset({"tasks:submit"}))
                ),
            )
        )
        client = TestClient(app)

        denied = client.post("/v1/tasks", json={"goal": "unscoped"})
        assert denied.status_code == 403
        allowed = client.post(
            "/v1/tasks",
            json={
                "goal": "scoped",
                "constraints": {
                    "target_host_id": "host-a",
                    "target_device_id": "device-a",
                },
            },
        )
        assert allowed.status_code == 201, allowed.text
    finally:
        database.close()


def test_governance_routes_persist_revisioned_policies(tmp_path) -> None:
    database = CloudDatabase(f"sqlite:///{(tmp_path / 'governance-api.sqlite3').as_posix()}")
    try:
        _add_user(database.repository)
        database.repository.upsert_host(
            "host-a", address=None, last_seen_at=datetime.now(UTC)
        )
        app = FastAPI()
        app.include_router(
            create_governance_router(
                repository=database.repository,
                auth_provider=_PrincipalProvider(
                    Principal(id="admin", scopes=frozenset({"governance:admin", "governance:read"}))
                ),
            )
        )
        client = TestClient(app)

        user_policy = client.put(
            "/v1/users/user-a/submission-policy",
            json={"submission_enabled": False, "allowed_host_ids": []},
        )
        assert user_policy.status_code == 200, user_policy.text
        assert user_policy.json()["revision"] == 1
        host_policy = client.put(
            "/v1/hosts/host-a/governance-policy",
            json={"max_active_tasks": 2, "daily_token_budget": 1000},
        )
        assert host_policy.status_code == 200, host_policy.text
        assert host_policy.json()["revision"] == 1
        reread = client.get("/v1/hosts/host-a/governance-policy")
        assert reread.json()["daily_token_budget"] == 1000
        conflict = client.put(
            "/v1/hosts/host-a/governance-policy",
            json={
                "max_active_tasks": 3,
                "daily_token_budget": 1000,
                "expected_revision": 1,
            },
        )
        assert conflict.status_code == 200
        stale = client.put(
            "/v1/hosts/host-a/governance-policy",
            json={"max_active_tasks": 4, "expected_revision": 1},
        )
        assert stale.status_code == 409
        usage = client.get("/v1/hosts/host-a/token-usage")
        assert usage.status_code == 200
        assert usage.json()["used_tokens"] == 0
        assert usage.json()["daily_token_budget"] == 1000
    finally:
        database.close()


def test_absent_policy_preserves_existing_submission_behavior() -> None:
    enforce_user_submission_policy(None, target_host_id=None, target_device_id=None)