litellm/tests/integration/database/test_roi_observed.py

import asyncio
from collections.abc import AsyncIterator
from datetime import date, datetime, timezone
from typing import Final
from urllib.parse import parse_qs, urlsplit

import httpx
import pytest
import pytest_asyncio
from fastapi import FastAPI, HTTPException
from pydantic import SecretStr

from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.roi_calculator.github import SourceError
from litellm.proxy.roi_calculator.oauth import (
    OAuthConfig,
    TokenGrant,
    begin_authorization,
    connected_settings,
    consume_state,
    exchange_code,
    save_grant,
)
from litellm.proxy.roi_calculator.observed_sync import ObservedSyncManager, Progress
from litellm.proxy.roi_calculator.settings import load_settings, load_stored_settings, save_settings
from litellm.proxy.roi_calculator.sync_store import SyncStore
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.config_repository import ConfigRepository
from litellm.types.roi_observed import ObservedData, ObservedPeriodData, ObservedWindow
from tests.integration._support.database import scratch_database, write_rows


@pytest_asyncio.fixture(loop_scope="function")
async def repository(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[ConfigRepository]:
    with scratch_database() as url:
        write_rows(
            'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB NOT NULL, '
            "last_run_at TIMESTAMP NOT NULL DEFAULT NOW(), reload_revision BIGINT NOT NULL DEFAULT 0)",
            (),
            database_url=url,
        )
        monkeypatch.setenv("DATABASE_URL", url)
        monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
        monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-observed-roi-salt-0123456789")
        client: Final = PrismaClient(url, ProxyLogging(UserApiKeyCache()))
        await client.connect()
        try:
            yield ConfigRepository(client, use_writer=True)
        finally:
            await client.disconnect()


def _config(provider: str = "github") -> OAuthConfig:
    from pydantic import TypeAdapter

    from litellm.proxy.roi_calculator.oauth import Provider

    selected: Final = TypeAdapter(Provider).validate_python(provider)
    return OAuthConfig(
        selected,
        "https://api.github.com" if provider == "github" else "https://gitlab.com/api/v4",
        f"https://{provider}.com",
        "test-client",
        SecretStr("test-client-secret"),
        "https://gateway.example.test",
    )


def _report() -> ObservedData:
    period: Final = ObservedPeriodData(
        window=ObservedWindow(start=date(2026, 9, 1), end=date(2026, 9, 28)), pulls=(), issues=(), spend=()
    )
    return ObservedData(
        source_provider="github",
        source_api_url="https://api.github.com",
        repos=("org/repo",),
        captured_at=datetime.now(timezone.utc),
        gateway_emails=(),
        current=period,
        previous=period,
        last_year=period,
    )


@pytest.mark.asyncio
async def test_settings_compare_and_swap_rejects_stale_writers(repository: ConfigRepository) -> None:
    assert await repository.set_param_if_revision("settings", {"revision": 1, "account": "ari"}, 0)
    assert not await repository.set_param_if_revision("settings", {"revision": 1, "account": "bea"}, 0)
    writes: Final = await asyncio.gather(
        *(repository.set_param_if_revision("settings", {"revision": 2, "account": name}, 1) for name in ("bea", "cam"))
    )
    assert sum(writes) == 1
    saved: Final = await repository.get_param("settings")
    assert saved is not None and saved.param_value in (
        {"revision": 2, "account": "bea"},
        {"revision": 2, "account": "cam"},
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("provider", ("github", "gitlab"))
async def test_authorization_uses_pkce_single_use_state_and_encrypted_credentials(
    repository: ConfigRepository, provider: str
) -> None:
    config: Final = _config(provider)
    url, nonce = await begin_authorization(repository, config)
    params: Final = parse_qs(urlsplit(url).query)
    assert params["code_challenge_method"] == ["S256"]
    assert params["redirect_uri"] == [config.redirect_uri]
    state: Final = await consume_state(repository, params["state"][0], nonce, config)
    assert state.verifier.get_secret_value() != "**********"

    def respond(request: httpx.Request) -> httpx.Response:
        body: Final = parse_qs(request.content.decode())
        assert body["code_verifier"] == [state.verifier.get_secret_value()]
        assert body["code"] == ["test-code"]
        assert body["client_secret"] == ["test-client-secret"]
        assert body["redirect_uri"] == [config.redirect_uri]
        return httpx.Response(
            200, json={"access_token": "test-access", "refresh_token": "test-refresh", "expires_in": 3600}
        )

    grant: Final = await exchange_code(config, state, "test-code", httpx.MockTransport(respond))
    await save_grant(repository, config, grant, revision=state.settings_revision)
    stored: Final = await load_stored_settings(repository)
    assert "test-access" not in stored.model_dump_json() and "test-refresh" not in stored.model_dump_json()
    connected: Final = await load_settings(repository)
    assert (
        connected.github_token if provider == "github" else connected.gitlab_token
    ).get_secret_value() == "test-access"
    assert connected.oauth_refresh_token.get_secret_value() == "test-refresh"
    with pytest.raises(HTTPException, match="already used"):
        await consume_state(repository, params["state"][0], nonce, config)


@pytest.mark.asyncio
async def test_authorization_rejects_a_different_browser_and_changed_settings(repository: ConfigRepository) -> None:
    config: Final = _config()
    url, nonce = await begin_authorization(repository, config)
    state: Final = parse_qs(urlsplit(url).query)["state"][0]
    with pytest.raises(HTTPException, match="same browser"):
        await consume_state(repository, state, "another-browser", config)
    new_url, new_nonce = await begin_authorization(repository, config)
    verified: Final = await consume_state(repository, parse_qs(urlsplit(new_url).query)["state"][0], new_nonce, config)
    await save_grant(repository, config, TokenGrant(access_token=SecretStr("first")))
    with pytest.raises(HTTPException, match="changed during authorization"):
        await save_grant(
            repository, config, TokenGrant(access_token=SecretStr("stale")), revision=verified.settings_revision
        )
    assert (await load_settings(repository)).github_token.get_secret_value() == "first"
    assert nonce != new_nonce


@pytest.mark.asyncio
async def test_concurrent_app_refresh_rotates_once_and_preserves_account_links(
    repository: ConfigRepository, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setenv("LITELLM_ROI_GITHUB_CLIENT_ID", "test-client")
    monkeypatch.setenv("LITELLM_ROI_GITHUB_CLIENT_SECRET", "test-client-secret")
    monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.test")
    await save_grant(
        repository,
        _config(),
        TokenGrant(access_token=SecretStr("old-access"), refresh_token=SecretStr("old-refresh"), expires_in=1),
    )
    stored: Final = await load_stored_settings(repository)
    settings: Final = (await load_settings(repository)).model_copy(update={"identity_map": {"ari": "ari@example.test"}})
    await save_settings(repository, settings, stored.github_token, stored.estimator_key, revision=stored.revision)
    requests: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue()

    def respond(request: httpx.Request) -> httpx.Response:
        requests.put_nowait(request)
        assert parse_qs(request.content.decode())["refresh_token"] == ["old-refresh"]
        return httpx.Response(
            200, json={"access_token": "new-access", "refresh_token": "new-refresh", "expires_in": 3600}
        )

    results: Final = await asyncio.gather(
        *(connected_settings(repository, httpx.MockTransport(respond)) for _ in range(12))
    )
    assert requests.qsize() == 1
    assert all(result.github_token.get_secret_value() == "new-access" for result in results)
    assert all(result.identity_map == {"ari": "ari@example.test"} for result in results)
    assert (await load_settings(repository)).oauth_refresh_token.get_secret_value() == "new-refresh"


@pytest.mark.asyncio
async def test_refresh_cannot_restore_a_connection_replaced_while_the_provider_responds(
    repository: ConfigRepository, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setenv("LITELLM_ROI_GITHUB_CLIENT_ID", "test-client")
    monkeypatch.setenv("LITELLM_ROI_GITHUB_CLIENT_SECRET", "test-client-secret")
    monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.test")
    await save_grant(
        repository,
        _config(),
        TokenGrant(access_token=SecretStr("old"), refresh_token=SecretStr("refresh"), expires_in=1),
    )
    started: Final = asyncio.Event()
    release: Final = asyncio.Event()

    async def respond(request: httpx.Request) -> httpx.Response:
        started.set()
        await release.wait()
        return httpx.Response(200, json={"access_token": "late-token", "expires_in": 3600})

    pending: Final = asyncio.create_task(connected_settings(repository, httpx.MockTransport(respond)))
    await asyncio.wait_for(started.wait(), 2)
    await save_grant(repository, _config(), TokenGrant(access_token=SecretStr("replacement"), expires_in=3600))
    release.set()
    result: Final = await asyncio.wait_for(pending, 2)
    assert result.github_token.get_secret_value() == "replacement"
    assert (await load_settings(repository)).github_token.get_secret_value() == "replacement"


@pytest.mark.asyncio
async def test_failed_cancelled_and_stale_workers_cannot_replace_the_published_report(
    repository: ConfigRepository,
) -> None:
    store: Final = SyncStore(repository.prisma_client, "roi_observed")
    manager: Final = ObservedSyncManager()
    report: Final = _report()
    await repository.set_param("roi_observed_report", report.model_dump(mode="json"))
    release: Final = asyncio.Event()

    async def build(progress: Progress) -> ObservedData:
        await release.wait()
        raise SourceError("source failure")

    assert await manager.start(build, store)
    assert not await ObservedSyncManager().start(build, store)
    await store.cancel()
    release.set()
    await manager.cancel()
    published: Final = await repository.get_param("roi_observed_report")
    assert published is not None and ObservedData.model_validate(published.param_value) == report

    async def complete(progress: Progress) -> ObservedData:
        return report

    assert await manager.start(complete, store)

    async def finished() -> None:
        for _ in range(200):
            status: Final = await store.status()
            if status is not None and not status.running:
                assert status.phase == "complete", status
                return
            await asyncio.sleep(0.01)
        pytest.fail("Observed report did not publish")

    await asyncio.wait_for(finished(), 5)
    saved: Final = await repository.get_param("roi_observed_report")
    assert saved is not None and ObservedData.model_validate(saved.param_value) == report


@pytest.mark.asyncio
@pytest.mark.parametrize("provider,install_first", (("github", False), ("github", True), ("gitlab", False)))
async def test_app_callback_round_trip_and_admin_authorization(
    repository: ConfigRepository, monkeypatch: pytest.MonkeyPatch, provider: str, install_first: bool
) -> None:
    from fastapi import FastAPI

    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
    from litellm.proxy.management_endpoints.roi_observed_endpoints import (
        get_oauth_repository,
        get_observed_transport,
        router,
    )
    from litellm.proxy.roi_calculator.settings import get_roi_config_repository

    config: Final = _config(provider)
    monkeypatch.setenv(f"LITELLM_ROI_{provider.upper()}_CLIENT_ID", config.client_id)
    monkeypatch.setenv(f"LITELLM_ROI_{provider.upper()}_CLIENT_SECRET", config.client_secret.get_secret_value())
    monkeypatch.setenv("PROXY_BASE_URL", config.proxy_url)
    monkeypatch.setenv("LITELLM_ROI_GITHUB_APP_SLUG", "example-roi")

    def respond(request: httpx.Request) -> httpx.Response:
        if request.method == "POST":
            assert str(request.url) == config.token_url
            if parse_qs(request.content.decode()).get("code") == ["rejected-code"]:
                return httpx.Response(400, json={"error": "invalid_grant"})
            return httpx.Response(
                200,
                json={
                    "access_token": "provider-test-access",
                    "refresh_token": "provider-test-refresh",
                    "expires_in": 3600,
                },
            )
        assert request.headers["Authorization"] == "Bearer provider-test-access"
        assert "PRIVATE-TOKEN" not in request.headers
        return httpx.Response(200, json=[])

    app: Final = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_roi_config_repository] = lambda: repository
    app.dependency_overrides[get_oauth_repository] = lambda: repository
    app.dependency_overrides[get_observed_transport] = lambda: httpx.MockTransport(respond)
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url=config.proxy_url) as client:
        for role, expected in (
            (LitellmUserRoles.INTERNAL_USER, 403),
            (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, 403),
            (LitellmUserRoles.PROXY_ADMIN, 200),
        ):
            app.dependency_overrides[user_api_key_auth] = lambda role=role: UserAPIKeyAuth(user_role=role)
            response: Final = await client.post(f"/roi-calculator/observed/oauth/{provider}/start")
            assert response.status_code == expected, response.text
        successful: Final = await client.post(
            f"/roi-calculator/observed/oauth/{provider}/start", params={"install": str(install_first).lower()}
        )
        assert "HttpOnly" in successful.headers["set-cookie"] and "Secure" in successful.headers["set-cookie"]
        initial_state: Final = parse_qs(urlsplit(successful.json()["url"]).query)["state"][0]
        installed: Final = (
            await client.get("/roi-calculator/observed/oauth/github/installed", params={"state": initial_state})
            if install_first
            else None
        )
        if install_first:
            assert urlsplit(successful.json()["url"]).path == "/apps/example-roi/installations/new"
            assert installed is not None
            assert installed.status_code == 303, installed.text
            assert urlsplit(installed.headers["location"]).path == "/login/oauth/authorize"
            assert parse_qs(urlsplit(installed.headers["location"]).query)["state"][0] != initial_state
        state: Final = (
            parse_qs(urlsplit(installed.headers["location"]).query)["state"][0] if installed else initial_state
        )
        callback: Final = await client.get(
            f"/roi-calculator/observed/oauth/{provider}/callback", params={"state": state, "code": "test-code"}
        )
        assert callback.status_code == 303, callback.text
        assert callback.headers["location"] == config.proxy_url + f"/ui/roi-calculator/?connected={provider}"
        saved: Final = await client.get("/roi-calculator/observed/settings")
        assert saved.json()["has_token"] is True and saved.json()["connection_type"] == "app"
        assert "provider-test-access" not in saved.text and "provider-test-refresh" not in saved.text
        replay: Final = await client.get(
            f"/roi-calculator/observed/oauth/{provider}/callback", params={"state": state, "code": "test-code"}
        )
        assert replay.status_code == 303
        assert replay.headers["location"] == config.proxy_url + "/ui/roi-calculator/?connection_failed=1"
        restart: Final = await client.post(f"/roi-calculator/observed/oauth/{provider}/start")
        denied_state: Final = parse_qs(urlsplit(restart.json()["url"]).query)["state"][0]
        denied: Final = await client.get(
            f"/roi-calculator/observed/oauth/{provider}/callback",
            params={"state": denied_state, "error": "access_denied"},
        )
        assert denied.status_code == 303
        assert denied.headers["location"] == config.proxy_url + "/ui/roi-calculator/?connection_cancelled=1"
        unchanged: Final = await client.get("/roi-calculator/observed/settings")
        assert unchanged.json() == saved.json()
        expired_installation: Final = (
            await client.get("/roi-calculator/observed/oauth/github/installed", params={"state": initial_state})
            if install_first
            else None
        )
        if expired_installation is not None:
            assert expired_installation.status_code == 303
            assert expired_installation.headers["location"].endswith("?connection_failed=1")
        retry: Final = await client.post(f"/roi-calculator/observed/oauth/{provider}/start")
        rejected: Final = await client.get(
            f"/roi-calculator/observed/oauth/{provider}/callback",
            params={"state": parse_qs(urlsplit(retry.json()["url"]).query)["state"][0], "code": "rejected-code"},
        )
        assert rejected.status_code == 303
        assert rejected.headers["location"] == config.proxy_url + "/ui/roi-calculator/?connection_failed=1"
        assert "litellm_roi_oauth" not in client.cookies
        assert (await client.get("/roi-calculator/observed/settings")).json() == saved.json()


@pytest.mark.asyncio
async def test_identity_api_combines_accounts_and_removes_automatic_links_without_resync(
    repository: ConfigRepository,
) -> None:
    from fastapi import FastAPI

    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
    from litellm.proxy.management_endpoints.roi_observed_endpoints import router
    from litellm.proxy.roi_calculator.settings import get_roi_config_repository
    from litellm.types.roi_observed import ObservedPull

    write_rows('CREATE TABLE "LiteLLM_UserTable" (user_id TEXT PRIMARY KEY, user_email TEXT)', ())
    write_rows(
        'INSERT INTO "LiteLLM_UserTable" VALUES (%s,%s),(%s,%s)', ("ari", "ari@example.test", "bea", "bea@example.test")
    )
    settings: Final = (await load_settings(repository)).model_copy(update={"repos": ("org/repo",)})
    await save_settings(repository, settings, "", "")
    base: Final = _report()
    pulls: Final = tuple(
        ObservedPull(
            repo="org/repo",
            number=index,
            title="Change",
            url=f"https://github.com/org/repo/pull/{index}",
            author=login,
            profile_email="ari@example.test" if login == "ari" else "",
            merged_at=base.captured_at,
        )
        for index, login in enumerate(("ari", "old-ari"))
    )
    data: Final = base.model_copy(
        update={
            "gateway_emails": ("ari@example.test", "bea@example.test"),
            "current": base.current.model_copy(update={"pulls": pulls}),
        }
    )
    await repository.set_param("roi_observed_report", data.model_dump(mode="json"))
    app: Final = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_roi_config_repository] = lambda: repository
    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
    async with httpx.AsyncClient(
        transport=httpx.ASGITransport(app=app), base_url="https://gateway.example.test"
    ) as client:
        linked: Final = await client.put(
            "/roi-calculator/observed/identities", json={"email": "ari@example.test", "logins": ["ari", "old-ari"]}
        )
        assert linked.status_code == 200, linked.text
        person: Final = linked.json()["report"]["people"][0]
        assert person["logins"] == ["ari", "old-ari"] and person["periods"]["current"]["merged_prs"] == 2
        conflict: Final = await client.put(
            "/roi-calculator/observed/identities", json={"email": "bea@example.test", "logins": ["old-ari"]}
        )
        assert conflict.status_code == 409
        unknown: Final = await client.put(
            "/roi-calculator/observed/identities", json={"email": "missing@example.test", "logins": []}
        )
        assert unknown.status_code == 422
        removed: Final = await client.put(
            "/roi-calculator/observed/identities", json={"email": "ari@example.test", "logins": []}
        )
        assert removed.json()["report"]["people"] == []
        assert [login.split(":", 1)[-1] for login in removed.json()["report"]["unmatched_logins"]] == ["ari", "old-ari"]
        app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
        )
        readonly: Final = await client.get("/roi-calculator/observed/report")
        assert readonly.status_code == 200
        denied: Final = await client.put(
            "/roi-calculator/observed/identities", json={"email": "ari@example.test", "logins": []}
        )
        assert denied.status_code == 403


@pytest.mark.asyncio
async def test_both_app_connections_keep_repositories_credentials_and_account_links(
    repository: ConfigRepository,
) -> None:
    from litellm.proxy.roi_calculator.settings import connection_id, stored_connections

    await save_grant(
        repository,
        _config(),
        TokenGrant(access_token=SecretStr("github-access"), refresh_token=SecretStr("github-refresh")),
    )
    stored: Final = await load_stored_settings(repository)
    github: Final = (await load_settings(repository)).model_copy(
        update={"repos": ("org/service", "org/docs"), "identity_map": {"ari": "ari@example.test"}}
    )
    await save_settings(repository, github, stored.github_token, stored.estimator_key, revision=stored.revision)
    await save_grant(
        repository,
        _config("gitlab"),
        TokenGrant(access_token=SecretStr("gitlab-access"), refresh_token=SecretStr("gitlab-refresh")),
    )
    gitlab_id: Final = connection_id("gitlab", _config("gitlab").api_url)
    github_id: Final = connection_id("github", _config().api_url)
    first: Final = await load_settings(repository, selected_id=github_id)
    second: Final = await load_settings(repository, selected_id=gitlab_id)
    assert first.repos == ("org/service", "org/docs")
    assert first.identity_map == {"ari": "ari@example.test"}
    assert first.github_token.get_secret_value() == "github-access"
    assert first.oauth_refresh_token.get_secret_value() == "github-refresh"
    assert second.gitlab_token.get_secret_value() == "gitlab-access"
    assert second.oauth_refresh_token.get_secret_value() == "gitlab-refresh"
    assert second.identity_map == {} and second.repos == ()
    await save_grant(repository, _config(), TokenGrant(access_token=SecretStr("github-rotated")), previous=first)
    preserved: Final = await load_settings(repository, selected_id=gitlab_id)
    assert preserved.model_dump(exclude={"github_token", "github_api_url"}) == second.model_dump(
        exclude={"github_token", "github_api_url"}
    )
    assert (await load_settings(repository, selected_id=github_id)).github_token.get_secret_value() == "github-rotated"
    persisted: Final = await load_stored_settings(repository)
    assert len(stored_connections(persisted)) == 2
    assert all(
        token not in persisted.model_dump_json()
        for token in ("github-access", "github-rotated", "gitlab-access", "github-refresh", "gitlab-refresh")
    )


@pytest.mark.asyncio
async def test_cross_provider_account_form_saves_atomically_without_legacy_logins(repository: ConfigRepository) -> None:
    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.management_endpoints.roi_observed_endpoints import router
    from litellm.proxy.roi_calculator.settings import get_roi_config_repository, stored_connections, write_admin

    write_rows('CREATE TABLE "LiteLLM_UserTable" (user_id TEXT PRIMARY KEY, user_email TEXT)', ())
    write_rows(
        'INSERT INTO "LiteLLM_UserTable" VALUES (%s, %s), (%s, %s)',
        ("ari", "ari@example.test", "bea", "bea@example.test"),
    )
    await save_grant(repository, _config("github"), TokenGrant(access_token=SecretStr("github-token")))
    await save_grant(repository, _config("gitlab"), TokenGrant(access_token=SecretStr("gitlab-token")))
    connections: Final = stored_connections(await load_stored_settings(repository))
    accounts: Final = tuple({"connection_id": entry.id, "login": "ari"} for entry in connections)
    app: Final = FastAPI()
    app.include_router(router)

    async def config_repository() -> ConfigRepository:
        return repository

    async def admin() -> UserAPIKeyAuth:
        return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)

    app.dependency_overrides[get_roi_config_repository] = config_repository
    app.dependency_overrides[write_admin] = admin
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://gateway.test") as client:
        saved: Final = await client.put(
            "/roi-calculator/observed/identities", json={"email": "ari@example.test", "accounts": accounts}
        )
        assert saved.status_code == 200, saved.text
        after: Final = stored_connections(await load_stored_settings(repository))
        assert all(entry.identity_map == {"ari": "ari@example.test"} for entry in after)
        conflicting: Final = ({"connection_id": connections[0].id, "login": "bea"}, accounts[1])
        rejected: Final = await client.put(
            "/roi-calculator/observed/identities", json={"email": "bea@example.test", "accounts": conflicting}
        )
        assert rejected.status_code == 409, rejected.text
        assert stored_connections(await load_stored_settings(repository)) == after


@pytest.mark.asyncio
async def test_http_sync_combines_providers_retains_period_and_recovers_invalid_reports(
    repository: ConfigRepository,
) -> None:
    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
    from litellm.proxy.management_endpoints.roi_observed_endpoints import (
        get_observed_manager,
        get_observed_transport,
        router,
    )
    from litellm.proxy.roi_calculator.settings import get_roi_config_repository
    from litellm.types.roi_observed import ObservedReportResponse, ObservedSettings

    write_rows('CREATE TABLE "LiteLLM_UserTable" (user_id TEXT PRIMARY KEY, user_email TEXT)', ())
    write_rows(
        'CREATE TABLE "LiteLLM_DailyUserSpend" (user_id TEXT, date TEXT, spend DOUBLE PRECISION, api_requests INTEGER)',
        (),
    )
    write_rows(
        'CREATE TABLE "LiteLLM_SpendLogs" (spend DOUBLE PRECISION, request_tags JSONB, metadata JSONB, "startTime" TIMESTAMP)',
        (),
    )

    def respond(request: httpx.Request) -> httpx.Response:
        if request.headers.get("Authorization") == "Bearer rejected":
            return httpx.Response(401, json={"message": "Bad credentials"})
        if request.url.path == "/graphql":
            return httpx.Response(
                200, json={"data": {"search": {"issueCount": 0, "nodes": [], "pageInfo": {"hasNextPage": False}}}}
            )
        if request.url.path.startswith("/repos/"):
            return httpx.Response(200, json={"full_name": "org/service", "has_issues": True})
        if request.url.path in ("/user/repos", "/api/v4/projects"):
            return httpx.Response(200, json=[])
        if request.url.path == "/api/v4/projects/org/service":
            return httpx.Response(200, json={"id": 1, "path_with_namespace": "org/service"})
        if request.url.path.endswith(("/merge_requests", "/issues")):
            return httpx.Response(200, json=[])
        raise AssertionError(str(request.url))

    manager: Final = ObservedSyncManager()
    app: Final = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_roi_config_repository] = lambda: repository
    app.dependency_overrides[get_observed_manager] = lambda: manager
    app.dependency_overrides[get_observed_transport] = lambda: httpx.MockTransport(respond)
    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)

    async def finish() -> None:
        for _ in range(200):
            status: Final = await SyncStore(repository.prisma_client, "roi_observed").status()
            if status and not status.running:
                assert status.phase == "complete", status.error
                return
            await asyncio.sleep(0.01)
        pytest.fail("Sync did not finish")

    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://gateway.test") as client:
        assert (await client.get("/roi-calculator/observed/report")).json() == {"report": None}
        assert (await client.post("/roi-calculator/observed/sync")).status_code == 409
        for provider, url in (("github", "https://api.github.com"), ("gitlab", "https://gitlab.com/api/v4")):
            result: Final = await client.put(
                "/roi-calculator/observed/settings",
                json={
                    "source_provider": provider,
                    "api_url": url,
                    "token": "test-token",
                    "repos": ["org/service"],
                    "update_interval_minutes": 0,
                },
            )
            assert result.status_code == 200, result.text
        settings: Final = ObservedSettings.model_validate(
            (await client.get("/roi-calculator/observed/settings")).json()
        )
        assert len(settings.connections) == 2 and all(entry.has_token for entry in settings.connections)
        for entry in settings.connections:
            repositories: Final = await client.get(
                "/roi-calculator/observed/repositories", params={"connection": entry.id}
            )
            assert repositories.status_code == 200, repositories.text
        rejected: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "source_provider": "github",
                "api_url": "https://api.github.com",
                "token": "rejected",
                "repos": ["org/service"],
            },
        )
        assert rejected.status_code == 502, rejected.text
        assert (
            ObservedSettings.model_validate((await client.get("/roi-calculator/observed/settings")).json()) == settings
        )
        for params in ({"days": 7}, {}):
            started: Final = await client.post("/roi-calculator/observed/sync", params=params)
            assert started.status_code == 202, started.text
            await asyncio.wait_for(finish(), 5)
            response: Final = await client.get("/roi-calculator/observed/report")
            report: Final = ObservedReportResponse.model_validate(response.json()).report
            assert report is not None and report.source_provider == "mixed"
            assert len(report.connections) == 2 and report.periods.current.merged_prs == 0
            assert all(
                (period.window.end - period.window.start).days == 6
                for period in (report.periods.current, report.periods.previous, report.periods.last_year)
            )
        await repository.set_param("roi_observed_report", {"invalid": True})
        assert (await client.get("/roi-calculator/observed/report")).status_code == 500
        recovered: Final = await client.post("/roi-calculator/observed/sync")
        assert recovered.status_code == 202, recovered.text
        await asyncio.wait_for(finish(), 5)
        rebuilt: Final = ObservedReportResponse.model_validate(
            (await client.get("/roi-calculator/observed/report")).json()
        ).report
        assert (
            rebuilt is not None
            and (rebuilt.periods.current.window.end - rebuilt.periods.current.window.start).days == 27
        )


@pytest.mark.asyncio
@pytest.mark.parametrize("provider", ("github", "gitlab"))
async def test_host_whitespace_preserves_app_credentials_and_account_matches(
    repository: ConfigRepository, provider: str
) -> None:
    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
    from litellm.proxy.management_endpoints.roi_observed_endpoints import get_observed_transport, router
    from litellm.proxy.roi_calculator.settings import connection_id, get_roi_config_repository

    config: Final = _config(provider)
    granted: Final = await save_grant(
        repository,
        config,
        TokenGrant(access_token=SecretStr("test-access"), refresh_token=SecretStr("test-refresh"), expires_in=3600),
    )
    stored: Final = await load_stored_settings(repository)
    before: Final = granted.model_copy(
        update={"repos": ("org/service",), "identity_map": {"ari": "ari@example.test"}, "ignored_logins": ("bea",)}
    )
    await save_settings(
        repository, before, stored.github_token, stored.estimator_key, stored.gitlab_token, revision=stored.revision
    )

    def respond(request: httpx.Request) -> httpx.Response:
        assert request.headers.get("Authorization") == "Bearer test-access"
        if "/repos/" in request.url.path:
            return httpx.Response(200, json={"full_name": "org/service"})
        if request.url.path.endswith("/merge_requests"):
            return httpx.Response(200, json=[])
        return httpx.Response(200, json={"id": 1, "path_with_namespace": "org/service"})

    app: Final = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_roi_config_repository] = lambda: repository
    app.dependency_overrides[get_observed_transport] = lambda: httpx.MockTransport(respond)
    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://gateway.test") as client:
        response: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "connection_id": connection_id(provider, config.api_url),
                "source_provider": provider,
                "api_url": f"  {config.api_url}/  ",
                "repos": ["org/service"],
            },
        )
        assert response.status_code == 200, response.text
    assert (await load_settings(repository)) == before


@pytest.mark.asyncio
async def test_connection_edits_replace_only_the_selected_host_and_keep_the_workspace_schedule(
    repository: ConfigRepository,
) -> None:
    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
    from litellm.proxy.management_endpoints.roi_observed_endpoints import get_observed_transport, router
    from litellm.proxy.roi_calculator.settings import get_roi_config_repository, stored_connections
    from litellm.types.roi_calculator import ROISettings
    from litellm.types.roi_observed import ObservedSettings

    legacy: Final = ROISettings(repos=("org/service",), estimator_model="legacy-model", update_interval_minutes=60)
    await save_settings(repository, legacy, "", "")

    def respond(request: httpx.Request) -> httpx.Response:
        if "/repos/" in request.url.path:
            return httpx.Response(200, json={"full_name": "org/service"})
        if "/projects/" in request.url.path:
            return httpx.Response(200, json={"id": 1, "path_with_namespace": "org/service"})
        raise AssertionError(str(request.url))

    app: Final = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_roi_config_repository] = lambda: repository
    app.dependency_overrides[get_observed_transport] = lambda: httpx.MockTransport(respond)
    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://gateway.test") as client:
        assert (await client.get("/roi-calculator/observed/settings")).status_code == 200
        assert (await load_settings(repository)) == legacy
        github: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "source_provider": "github",
                "api_url": "https://api.github.com",
                "token": "github-token",
                "repos": ["org/service"],
                "update_interval_minutes": 30,
            },
        )
        assert github.status_code == 200, github.text
        gitlab: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "source_provider": "gitlab",
                "api_url": "https://gitlab.com/api/v4",
                "token": "gitlab-token",
                "repos": ["org/service"],
            },
        )
        assert gitlab.status_code == 200, gitlab.text
        before: Final = stored_connections(await load_stored_settings(repository))
        edited: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "connection_id": github.json()["id"],
                "source_provider": "github",
                "api_url": "https://git.example.test/api/v3",
                "token": "enterprise-token",
                "repos": ["org/service"],
            },
        )
        assert edited.status_code == 200, edited.text
        settings: Final = ObservedSettings.model_validate(edited.json())
        assert len(settings.connections) == 2
        assert {entry.api_url for entry in settings.connections} == {
            "https://git.example.test/api/v3",
            "https://gitlab.com/api/v4",
        }
        assert all(entry.update_interval_minutes == 30 for entry in settings.connections)
        after: Final = stored_connections(await load_stored_settings(repository))
        assert next(entry for entry in after if entry.source_provider == "gitlab") == next(
            entry for entry in before if entry.source_provider == "gitlab"
        )
        assert (await load_settings(repository)).report_mode == "observed"
        stale: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "connection_id": github.json()["id"],
                "source_provider": "github",
                "api_url": "https://api.github.com",
                "repos": [],
            },
        )
        assert stale.status_code == 404, stale.text
        duplicate: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "connection_id": settings.id,
                "source_provider": "gitlab",
                "api_url": "https://gitlab.com/api/v4",
                "repos": [],
            },
        )
        assert duplicate.status_code == 409, duplicate.text
        assert stored_connections(await load_stored_settings(repository)) == after
        manual: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "source_provider": "gitlab",
                "api_url": "https://gitlab.com/api/v4",
                "repos": ["org/service"],
                "update_interval_minutes": 0,
            },
        )
        assert manual.status_code == 200, manual.text
        reselected: Final = await client.put(
            "/roi-calculator/observed/settings",
            json={
                "connection_id": settings.id,
                "source_provider": "github",
                "api_url": "https://git.example.test/api/v3",
                "repos": ["org/service"],
            },
        )
        assert reselected.status_code == 200, reselected.text
        assert all(
            entry.update_interval_minutes == 0
            for entry in ObservedSettings.model_validate(reselected.json()).connections
        )