Aether/tests/services/test_provider_auth_vertex_service_account.py

from __future__ import annotations

import json
import sys
import types
from types import SimpleNamespace

import pytest

from src.core.vertex_auth import VertexAuthService
from src.services.provider.auth import get_provider_auth


class _FakeQuery:
    def __init__(self, row: object | None) -> None:
        self._row = row

    def filter(self, *_args: object, **_kwargs: object) -> "_FakeQuery":
        return self

    def first(self) -> object | None:
        return self._row


class _FakeSessionCtx:
    def __init__(self, row: object | None) -> None:
        self._row = row

    def __enter__(self) -> "_FakeSessionCtx":
        return self

    def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
        _ = exc_type, exc, tb
        return False

    def query(self, _model: object) -> _FakeQuery:
        return _FakeQuery(self._row)


@pytest.mark.asyncio
async def test_get_provider_auth_vertex_service_account_uses_provider_proxy(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    sa_json = {
        "client_email": "svc@example.iam.gserviceaccount.com",
        "private_key": "-----BEGIN PRIVATE KEY-----\nTEST\n-----END PRIVATE KEY-----\n",
        "project_id": "demo-project",
    }
    provider_proxy = {"node_id": "provider-node", "enabled": True}
    provider = SimpleNamespace(proxy=provider_proxy)
    endpoint = SimpleNamespace(provider=provider)
    key = SimpleNamespace(
        auth_type="service_account",
        auth_config="enc_cfg",
        api_key="enc_key",
        provider=provider,
        proxy=None,
    )

    captured: dict[str, object] = {}

    monkeypatch.setattr(
        "src.core.crypto.crypto_service.decrypt",
        lambda value: json.dumps(sa_json) if value == "enc_cfg" else "",
    )

    def _fake_build_proxy_client_kwargs(
        proxy_config: dict[str, object] | None = None,
        *,
        timeout: float = 30.0,
        **_: object,
    ) -> dict[str, object]:
        captured["proxy_config"] = proxy_config
        captured["timeout"] = timeout
        return {"timeout": timeout}

    async def _fake_get_access_token(
        self: VertexAuthService,
        *,
        httpx_client_kwargs: dict[str, object] | None = None,
    ) -> str:
        captured["httpx_client_kwargs"] = httpx_client_kwargs
        return "ya29.test-token"

    monkeypatch.setattr(
        "src.services.proxy_node.resolver.build_proxy_client_kwargs",
        _fake_build_proxy_client_kwargs,
    )
    monkeypatch.setattr(VertexAuthService, "get_access_token", _fake_get_access_token)

    auth = await get_provider_auth(endpoint, key)  # type: ignore[arg-type]

    assert auth is not None
    assert captured["proxy_config"] == provider_proxy
    assert captured["timeout"] == 30
    assert captured["httpx_client_kwargs"] == {"timeout": 30}
    assert auth.auth_header == "Authorization"
    assert auth.auth_value == "Bearer ya29.test-token"
    assert auth.decrypted_auth_config == sa_json


@pytest.mark.asyncio
async def test_get_provider_auth_vertex_service_account_prefers_key_proxy(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    sa_json = {
        "client_email": "svc@example.iam.gserviceaccount.com",
        "private_key": "-----BEGIN PRIVATE KEY-----\nTEST\n-----END PRIVATE KEY-----\n",
        "project_id": "demo-project",
    }
    provider = SimpleNamespace(proxy={"node_id": "provider-node", "enabled": True})
    endpoint = SimpleNamespace(provider=provider)
    key_proxy = {"node_id": "key-node", "enabled": True}
    key = SimpleNamespace(
        auth_type="service_account",
        auth_config="enc_cfg",
        api_key="enc_key",
        provider=provider,
        proxy=key_proxy,
    )

    captured: dict[str, object] = {}

    monkeypatch.setattr(
        "src.core.crypto.crypto_service.decrypt",
        lambda value: json.dumps(sa_json) if value == "enc_cfg" else "",
    )

    def _fake_build_proxy_client_kwargs(
        proxy_config: dict[str, object] | None = None,
        *,
        timeout: float = 30.0,
        **_: object,
    ) -> dict[str, object]:
        captured["proxy_config"] = proxy_config
        return {"timeout": timeout}

    async def _fake_get_access_token(
        self: VertexAuthService,
        *,
        httpx_client_kwargs: dict[str, object] | None = None,
    ) -> str:
        return "ya29.test-token"

    monkeypatch.setattr(
        "src.services.proxy_node.resolver.build_proxy_client_kwargs",
        _fake_build_proxy_client_kwargs,
    )
    monkeypatch.setattr(VertexAuthService, "get_access_token", _fake_get_access_token)

    auth = await get_provider_auth(endpoint, key)  # type: ignore[arg-type]

    assert auth is not None
    assert captured["proxy_config"] == key_proxy


@pytest.mark.asyncio
async def test_get_provider_auth_vertex_service_account_uses_provider_id_lookup_without_touching_endpoint_provider(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    sa_json = {
        "client_email": "svc@example.iam.gserviceaccount.com",
        "private_key": "-----BEGIN PRIVATE KEY-----\nTEST\n-----END PRIVATE KEY-----\n",
        "project_id": "demo-project",
    }
    provider_proxy = {"node_id": "provider-node", "enabled": True}

    class _DetachedEndpoint:
        provider_id = "provider-1"

        @property
        def provider(self) -> object:
            raise RuntimeError("detached endpoint provider should not be lazy-loaded")

    fake_provider = SimpleNamespace(
        id="provider-1", provider_type="vertex_ai", proxy=provider_proxy
    )
    fake_database = types.ModuleType("src.database")
    fake_database.create_session = lambda: _FakeSessionCtx(fake_provider)
    fake_models = types.ModuleType("src.models.database")
    fake_models.Provider = type("Provider", (), {"id": "id"})

    monkeypatch.setitem(sys.modules, "src.database", fake_database)
    monkeypatch.setitem(sys.modules, "src.models.database", fake_models)
    monkeypatch.setattr(
        "src.core.crypto.crypto_service.decrypt",
        lambda value: json.dumps(sa_json) if value == "enc_cfg" else "",
    )

    captured: dict[str, object] = {}

    def _fake_build_proxy_client_kwargs(
        proxy_config: dict[str, object] | None = None,
        *,
        timeout: float = 30.0,
        **_: object,
    ) -> dict[str, object]:
        captured["proxy_config"] = proxy_config
        return {"timeout": timeout}

    async def _fake_get_access_token(
        self: VertexAuthService,
        *,
        httpx_client_kwargs: dict[str, object] | None = None,
    ) -> str:
        captured["httpx_client_kwargs"] = httpx_client_kwargs
        return "ya29.test-token"

    monkeypatch.setattr(
        "src.services.proxy_node.resolver.build_proxy_client_kwargs",
        _fake_build_proxy_client_kwargs,
    )
    monkeypatch.setattr(VertexAuthService, "get_access_token", _fake_get_access_token)

    key = SimpleNamespace(
        auth_type="service_account",
        auth_config="enc_cfg",
        api_key="enc_key",
        provider_id="provider-1",
        proxy=None,
    )

    auth = await get_provider_auth(_DetachedEndpoint(), key)  # type: ignore[arg-type]

    assert auth is not None
    assert captured["proxy_config"] == provider_proxy
    assert captured["httpx_client_kwargs"] == {"timeout": 30}