Pipeline/backend/tests/test_provider_credentials_usage_api.py

# ruff: noqa: INP001
"""Integration tests for provider credential live-usage API behavior."""

from __future__ import annotations

from uuid import uuid4

import pytest
from fastapi import APIRouter, FastAPI
from httpx import ASGITransport, AsyncClient
from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession

from app.api.deps import require_org_admin, require_org_member
from app.api.provider_credentials import router as provider_credentials_router
from app.core.time import utcnow
from app.db.session import get_session
from app.models.organization_members import OrganizationMember
from app.models.organizations import Organization
from app.models.provider_credentials import ProviderCredential
from app.services.organizations import OrganizationContext
from app.services.provider_usage import ProviderUsageLive


async def _make_engine() -> AsyncEngine:
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)
    return engine


def _build_test_app(
    session_maker: async_sessionmaker[AsyncSession],
    ctx: OrganizationContext,
) -> FastAPI:
    app = FastAPI()
    api_v1 = APIRouter(prefix="/api/v1")
    api_v1.include_router(provider_credentials_router)
    app.include_router(api_v1)

    async def _override_get_session() -> AsyncSession:
        async with session_maker() as session:
            yield session

    async def _override_require_org_member() -> OrganizationContext:
        return ctx

    async def _override_require_org_admin() -> OrganizationContext:
        return ctx

    app.dependency_overrides[get_session] = _override_get_session
    app.dependency_overrides[require_org_member] = _override_require_org_member
    app.dependency_overrides[require_org_admin] = _override_require_org_admin
    return app


@pytest.mark.asyncio
async def test_usage_response_includes_rate_limit_header_names(monkeypatch: pytest.MonkeyPatch) -> None:
    engine = await _make_engine()
    session_maker = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    organization = Organization(id=uuid4(), name="Pipeline")
    member = OrganizationMember(
        id=uuid4(),
        organization_id=organization.id,
        user_id=uuid4(),
        role="owner",
    )
    app = _build_test_app(
        session_maker,
        OrganizationContext(organization=organization, member=member),
    )

    credential = ProviderCredential(
        id=uuid4(),
        organization_id=organization.id,
        provider="anthropic",
        account_key="Claude",
        display_name="Claude",
        api_key="sk-ant-test",
        api_key_last_four="test",
        active=True,
    )

    async def _fake_fetch_provider_usage(**_: object) -> ProviderUsageLive:
        result = ProviderUsageLive(
            provider="anthropic",
            account_key="Claude",
            checked_at=utcnow(),
            reachable=True,
        )
        result.raw_headers = {
            "anthropic-ratelimit-requests-limit": "1000",
            "anthropic-ratelimit-requests-remaining": "999",
            "anthropic-ratelimit-requests-reset": "2026-05-21T12:00:00Z",
        }
        return result

    monkeypatch.setattr(
        "app.api.provider_credentials.fetch_provider_usage",
        _fake_fetch_provider_usage,
    )

    try:
        async with session_maker() as session:
            session.add(organization)
            session.add(credential)
            await session.commit()

        async with AsyncClient(
            transport=ASGITransport(app=app),
            base_url="http://testserver",
        ) as client:
            response = await client.get(f"/api/v1/provider-credentials/{credential.id}/usage")

        assert response.status_code == 200
        data = response.json()
        assert data["provider"] == "anthropic"
        assert data["reachable"] is True
        assert data["debug_rate_limit_headers"] == [
            "anthropic-ratelimit-requests-limit",
            "anthropic-ratelimit-requests-remaining",
            "anthropic-ratelimit-requests-reset",
        ]
    finally:
        await engine.dispose()


@pytest.mark.asyncio
async def test_test_endpoint_returns_live_result(monkeypatch: pytest.MonkeyPatch) -> None:
    engine = await _make_engine()
    session_maker = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    organization = Organization(id=uuid4(), name="Pipeline")
    member = OrganizationMember(
        id=uuid4(),
        organization_id=organization.id,
        user_id=uuid4(),
        role="owner",
    )
    app = _build_test_app(
        session_maker,
        OrganizationContext(organization=organization, member=member),
    )

    async def _fake_fetch_provider_usage(**kwargs: object) -> ProviderUsageLive:
        result = ProviderUsageLive(
            provider=str(kwargs["provider"]),
            account_key=str(kwargs["account_key"]),
            checked_at=utcnow(),
            reachable=True,
        )
        result.models = ["claude-sonnet-4-6"]
        result.raw_headers = {
            "anthropic-ratelimit-tokens-limit": "100000",
        }
        return result

    monkeypatch.setattr(
        "app.api.provider_credentials.fetch_provider_usage",
        _fake_fetch_provider_usage,
    )

    try:
        async with session_maker() as session:
            session.add(organization)
            await session.commit()

        async with AsyncClient(
            transport=ASGITransport(app=app),
            base_url="http://testserver",
        ) as client:
            response = await client.post(
                "/api/v1/provider-credentials/test",
                json={
                    "provider": "anthropic",
                    "account_key": "Claude",
                    "api_key": "sk-ant-test",
                },
            )

        assert response.status_code == 200
        data = response.json()
        assert data["provider"] == "anthropic"
        assert data["account_key"] == "Claude"
        assert data["reachable"] is True
        assert data["models"] == ["claude-sonnet-4-6"]
        assert data["debug_rate_limit_headers"] == ["anthropic-ratelimit-tokens-limit"]
    finally:
        await engine.dispose()