From 152e23372255f62e1114966223efe9ce712ec513 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 20 Jul 2026 16:50:23 -0700 Subject: [PATCH] feat(ptu): fetch azure_billing flat cost via Cost Management API Stage 6 POC on top of the LIT-1697 stack. Adds an azure_billing branch to the daily rollup so a reservation with cost_source=azure_billing pulls accrued cost from the Azure Cost Management API instead of computing the manual (ptu_count * cost_per_ptu) / days-in-month formula. The client is a thin async wrapper around one endpoint the rollup needs; Entra ID auth is inherited from get_azure_ad_token_from_entra_id (no new dependency). Deliberately narrow: no retry, no reconciliation for Azure's 24-72h reporting lag, no UI toggle for cost_source. Those are follow-ups. Gated behind a second UI settings flag enable_azure_ptu_billing_pull that composes below enable_ptu_cost_attribution. The scheduler builds a fetcher just-in-time so runtime flag flips take effect without a proxy restart; manual reservations keep working when either flag is off. Behavior changes: azure_billing cost_source is now accepted by /ptu_reservation/new when azure_resource_id is present. The rollup uses the injected fetcher for those; when no fetcher is available (flag off or env creds missing) the reservation is skipped with a warning and no sentinel row is written. --- .../azure_cost_management/__init__.py | 6 + .../azure_cost_management_client.py | 191 ++++++++++++++++++ litellm/proxy/dev_config.yaml | 2 + .../ptu_reservation_endpoints.py | 4 +- litellm/proxy/proxy_server.py | 41 +++- .../spend_tracking/ptu_reservation_rollup.py | 79 ++++++-- .../proxy_setting_endpoints.py | 7 + scripts/ptu_reservation_backfill.py | 33 ++- .../azure_cost_management/__init__.py | 0 .../test_azure_cost_management_client.py | 145 +++++++++++++ .../test_ptu_reservation_rollup.py | 98 +++++++-- 11 files changed, 573 insertions(+), 33 deletions(-) create mode 100644 litellm/integrations/azure_cost_management/__init__.py create mode 100644 litellm/integrations/azure_cost_management/azure_cost_management_client.py create mode 100644 tests/test_litellm/integrations/azure_cost_management/__init__.py create mode 100644 tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py diff --git a/litellm/integrations/azure_cost_management/__init__.py b/litellm/integrations/azure_cost_management/__init__.py new file mode 100644 index 00000000000..5eae86d6c4c --- /dev/null +++ b/litellm/integrations/azure_cost_management/__init__.py @@ -0,0 +1,6 @@ +from litellm.integrations.azure_cost_management.azure_cost_management_client import ( + AzureCostManagementClient, + AzureCostManagementError, +) + +__all__ = ["AzureCostManagementClient", "AzureCostManagementError"] diff --git a/litellm/integrations/azure_cost_management/azure_cost_management_client.py b/litellm/integrations/azure_cost_management/azure_cost_management_client.py new file mode 100644 index 00000000000..c7e778fd227 --- /dev/null +++ b/litellm/integrations/azure_cost_management/azure_cost_management_client.py @@ -0,0 +1,191 @@ +"""Thin async client for Azure Cost Management REST API. + +Used by the PTU reservation rollup to fetch billed cost for a specific +Azure resource id on a given day (LIT-4077). Deliberately narrow: exposes +one public method the rollup needs, not a general Cost Management SDK. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import date, datetime, time, timedelta, timezone +from typing import Any, Callable, Optional + +import httpx + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + +class AzureCostManagementError(Exception): + """Raised when the Cost Management API returns a non-2xx or an unexpected payload.""" + + +@dataclass(frozen=True, slots=True) +class AzureCostManagementConfig: + subscription_id: str + tenant_id: str + client_id: str + client_secret: str + api_version: str = "2023-11-01" + + @classmethod + def from_env(cls, subscription_id: str) -> AzureCostManagementConfig: + tenant_id = os.getenv("AZURE_TENANT_ID") + client_id = os.getenv("AZURE_CLIENT_ID") + client_secret = os.getenv("AZURE_CLIENT_SECRET") + missing = [ + name + for name, value in ( + ("AZURE_TENANT_ID", tenant_id), + ("AZURE_CLIENT_ID", client_id), + ("AZURE_CLIENT_SECRET", client_secret), + ) + if not value + ] + if missing: + raise AzureCostManagementError(f"Missing required env vars for Azure Cost Management auth: {missing}") + return cls( + subscription_id=subscription_id, + tenant_id=tenant_id or "", + client_id=client_id or "", + client_secret=client_secret or "", + ) + + +TokenProvider = Callable[[], str] + + +def _default_token_provider_factory(config: AzureCostManagementConfig) -> TokenProvider: + from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id + + return get_azure_ad_token_from_entra_id( + tenant_id=config.tenant_id, + client_id=config.client_id, + client_secret=config.client_secret, + scope="https://management.azure.com/.default", + ) + + +class AzureCostManagementClient: + """Fetch billed cost for a single Azure resource on a specific day. + + Auth: Entra ID service principal via env vars (``AZURE_TENANT_ID``, + ``AZURE_CLIENT_ID``, ``AZURE_CLIENT_SECRET``). The token provider is + dependency-injectable for tests. + + Currency: the API returns cost in the subscription's billing currency, + exposed on the response. Callers that require USD MUST inspect + ``last_currency`` after each call and handle non-USD. + """ + + def __init__( + self, + config: AzureCostManagementConfig, + *, + http_handler: Optional[AsyncHTTPHandler] = None, + token_provider: Optional[TokenProvider] = None, + ) -> None: + self._config = config + self._http = http_handler or AsyncHTTPHandler() + self._token_provider = token_provider or _default_token_provider_factory(config) + self._last_currency: Optional[str] = None + + @property + def last_currency(self) -> Optional[str]: + return self._last_currency + + async def get_daily_cost(self, resource_id: str, day: date) -> float: + """Return billed cost for ``resource_id`` on the UTC calendar day ``day``. + + Raises ``AzureCostManagementError`` on non-2xx responses or unexpected + payloads. Returns 0.0 when Azure reports no rows for the window (a + valid response, typically due to reporting lag or zero utilization). + """ + url = ( + "https://management.azure.com/subscriptions/" + f"{self._config.subscription_id}" + "/providers/Microsoft.CostManagement/query" + ) + start = datetime.combine(day, time.min, tzinfo=timezone.utc) + end = start + timedelta(days=1) - timedelta(microseconds=1) + body = { + "type": "ActualCost", + "timeframe": "Custom", + "timePeriod": { + "from": start.isoformat().replace("+00:00", "Z"), + "to": end.isoformat().replace("+00:00", "Z"), + }, + "dataset": { + "granularity": "Daily", + "aggregation": {"totalCost": {"name": "Cost", "function": "Sum"}}, + "filter": { + "dimensions": { + "name": "ResourceId", + "operator": "In", + "values": [resource_id], + } + }, + }, + } + token = self._token_provider() + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + } + try: + response = await self._http.post( + url=url, + json=body, + params={"api-version": self._config.api_version}, + headers=headers, + ) + except httpx.HTTPStatusError as exc: + raise AzureCostManagementError( + f"Azure Cost Management HTTP {exc.response.status_code}: {exc.response.text}" + ) from exc + return self._parse_cost_and_currency(response.json()) + + def _parse_cost_and_currency(self, payload: Any) -> float: + try: + properties = payload.get("properties", {}) if isinstance(payload, dict) else {} + rows = properties.get("rows") or [] + columns = properties.get("columns") or [] + except AttributeError as exc: + raise AzureCostManagementError(f"Unexpected payload shape: {payload!r}") from exc + + if not rows: + self._last_currency = None + return 0.0 + + column_index = {col.get("name"): i for i, col in enumerate(columns) if isinstance(col, dict)} + cost_idx: Optional[int] = None + for candidate in ("Cost", "PreTaxCost", "CostUSD"): + if candidate in column_index: + cost_idx = column_index[candidate] + break + currency_idx = column_index.get("Currency") + if cost_idx is None: + raise AzureCostManagementError(f"No Cost column in response; columns={list(column_index)}") + + try: + total = sum(float(row[cost_idx]) for row in rows) + except (TypeError, ValueError, IndexError) as exc: + raise AzureCostManagementError(f"Non-numeric cost value in rows: {rows}") from exc + + if currency_idx is not None and rows: + try: + self._last_currency = str(rows[0][currency_idx]) + except (IndexError, TypeError): + self._last_currency = None + else: + self._last_currency = None + + if self._last_currency is not None and self._last_currency.upper() != "USD": + verbose_logger.warning( + "Azure Cost Management returned currency=%s for resource %s; value written as-is without conversion", + self._last_currency, + "(resource redacted)", + ) + return total diff --git a/litellm/proxy/dev_config.yaml b/litellm/proxy/dev_config.yaml index f78431f694b..e88ad324a10 100644 --- a/litellm/proxy/dev_config.yaml +++ b/litellm/proxy/dev_config.yaml @@ -203,6 +203,8 @@ general_settings: # gs:// (Vertex) or s3:// (Bedrock) input_file_id. Requires a matching deployment # configured for the batched model. Defaults to false. # track_unmanaged_batch_cost: true + azure_ptu_billing: + subscription_id: os.environ/AZURE_SUBSCRIPTION_ID sandbox_tools: - sandbox_tool_name: e2b_sandbox diff --git a/litellm/proxy/management_endpoints/ptu_reservation_endpoints.py b/litellm/proxy/management_endpoints/ptu_reservation_endpoints.py index c40289a13f2..9d42578f9d8 100644 --- a/litellm/proxy/management_endpoints/ptu_reservation_endpoints.py +++ b/litellm/proxy/management_endpoints/ptu_reservation_endpoints.py @@ -90,10 +90,10 @@ async def new_ptu_reservation( _require_proxy_admin(user_api_key_dict) prisma_client = _require_db() - if body.cost_source == "azure_billing": + if body.cost_source == "azure_billing" and not body.azure_resource_id: raise HTTPException( status_code=400, - detail={"error": "cost_source='azure_billing' is not supported in this release"}, + detail={"error": "azure_resource_id is required when cost_source='azure_billing'"}, ) try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 085ea7342fd..07e810ff648 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7464,6 +7464,40 @@ def giveup(e): return result +def _build_azure_cost_fetcher_if_enabled() -> Optional[Any]: + """Return an AzureCostManagementClient when the pull flag + config + creds are all present, else None. + + Evaluated at rollup call time so runtime flag toggles take effect without + a proxy restart. + """ + if not general_settings.get("enable_azure_ptu_billing_pull", False): + return None + azure_ptu_billing = general_settings.get("azure_ptu_billing") or {} + subscription_id = azure_ptu_billing.get("subscription_id") + if not subscription_id: + verbose_proxy_logger.warning( + "enable_azure_ptu_billing_pull is true but general_settings.azure_ptu_billing.subscription_id is not set" + ) + return None + try: + from litellm.integrations.azure_cost_management import ( + AzureCostManagementClient, + AzureCostManagementError, + ) + from litellm.integrations.azure_cost_management.azure_cost_management_client import ( + AzureCostManagementConfig, + ) + + config = AzureCostManagementConfig.from_env(subscription_id=subscription_id) + return AzureCostManagementClient(config=config) + except Exception as exc: # noqa: BLE001 # missing creds/env; log once per invocation and skip azure_billing this run + verbose_proxy_logger.warning( + "Azure Cost Management client could not be built: %s. azure_billing reservations will no-op this run.", + exc, + ) + return None + + class ProxyStartupEvent: @classmethod def _initialize_startup_logging( @@ -7935,12 +7969,15 @@ class ProxyStartupEvent: run_ptu_reservation_rollup, ) + async def _scheduled_ptu_rollup() -> None: + azure_fetcher = _build_azure_cost_fetcher_if_enabled() + await run_ptu_reservation_rollup(prisma_client, azure_fetcher=azure_fetcher) + scheduler.add_job( - run_ptu_reservation_rollup, + _scheduled_ptu_rollup, "cron", hour=0, minute=15, - args=[prisma_client], id=PTU_ROLLUP_JOB_ID, replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, diff --git a/litellm/proxy/spend_tracking/ptu_reservation_rollup.py b/litellm/proxy/spend_tracking/ptu_reservation_rollup.py index 69de5b157a4..c1d3986b5e8 100644 --- a/litellm/proxy/spend_tracking/ptu_reservation_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_reservation_rollup.py @@ -9,13 +9,22 @@ existing unique constraint. from calendar import monthrange from dataclasses import dataclass from datetime import date, datetime, time, timedelta, timezone -from typing import Any +from typing import Any, Optional, Protocol from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_ROLLUP_JOB_ID, PTU_SENTINEL_API_KEY from litellm.repositories.ptu_reservation_repository import PTUReservationRepository +class AzureCostFetcher(Protocol): + """Dependency injected into the rollup for azure_billing reservations. + + Implemented by ``AzureCostManagementClient`` in production, mocked in tests. + """ + + async def get_daily_cost(self, resource_id: str, day: date) -> float: ... + + @dataclass(frozen=True, slots=True) class RollupResult: day: date @@ -28,14 +37,55 @@ def _days_in_month(day: date) -> int: return monthrange(day.year, day.month)[1] -def _compute_daily_flat_cost(reservation: Any, day: date) -> float: - """Return the flat cost attributable to ``day`` for a single reservation.""" - if reservation.cost_source != "manual": - return 0.0 - if reservation.ptu_count is None or reservation.cost_per_ptu is None: - return 0.0 - monthly_total = float(reservation.ptu_count) * float(reservation.cost_per_ptu) - return monthly_total / float(_days_in_month(day)) +async def _compute_daily_flat_cost( + reservation: Any, + day: date, + *, + azure_fetcher: Optional[AzureCostFetcher] = None, +) -> float: + """Return the flat cost attributable to ``day`` for a single reservation. + + manual: prorated (ptu_count * cost_per_ptu) / days_in_month. + azure_billing: live fetch via ``azure_fetcher`` when provided; 0.0 when + the pull is not configured (rollup logs a skip). + """ + if reservation.cost_source == "manual": + if reservation.ptu_count is None or reservation.cost_per_ptu is None: + return 0.0 + monthly_total = float(reservation.ptu_count) * float(reservation.cost_per_ptu) + return monthly_total / float(_days_in_month(day)) + if reservation.cost_source == "azure_billing": + if azure_fetcher is None or reservation.azure_resource_id is None: + verbose_proxy_logger.warning( + "PTU rollup: reservation=%s cost_source=azure_billing skipped " + "(enable_azure_ptu_billing_pull off or azure_resource_id missing)", + getattr(reservation, "id", "?"), + ) + return 0.0 + try: + fetched = await azure_fetcher.get_daily_cost(reservation.azure_resource_id, day) + verbose_proxy_logger.info( + "PTU rollup: azure_billing reservation=%s day=%s resource=%s returned $%.4f", + getattr(reservation, "id", "?"), + day.isoformat(), + reservation.azure_resource_id, + fetched, + ) + return fetched + except Exception as exc: # noqa: BLE001 # log and continue; one bad reservation must not stop the batch + verbose_proxy_logger.error( + "PTU rollup: azure fetch failed for reservation=%s day=%s: %s", + getattr(reservation, "id", "?"), + day.isoformat(), + exc, + ) + return 0.0 + verbose_proxy_logger.warning( + "PTU rollup: unknown cost_source=%s on reservation=%s; skipping", + reservation.cost_source, + getattr(reservation, "id", "?"), + ) + return 0.0 async def _upsert_ptu_daily_row( @@ -88,12 +138,16 @@ async def run_ptu_reservation_rollup( target_date: date | None = None, *, force: bool = False, + azure_fetcher: Optional[AzureCostFetcher] = None, ) -> RollupResult: """Rollup one UTC day of flat PTU cost across all active reservations. Defaults to yesterday UTC. ``force=True`` bypasses the feature-flag check - so the CLI backfill can run when the scheduler is off. Idempotent under - the LiteLLM_DailyTeamSpend unique constraint on every invocation path. + so the CLI backfill can run when the scheduler is off. ``azure_fetcher`` + is injected by the proxy startup wiring when ``enable_azure_ptu_billing_pull`` + is on; azure_billing reservations no-op with a warning when it is None. + Idempotent under the LiteLLM_DailyTeamSpend unique constraint on every + invocation path. """ if not force: from litellm.proxy.proxy_server import general_settings @@ -124,7 +178,7 @@ async def run_ptu_reservation_rollup( rows_written = 0 for reservation in reservations: - flat_cost = _compute_daily_flat_cost(reservation, day) + flat_cost = await _compute_daily_flat_cost(reservation, day, azure_fetcher=azure_fetcher) if flat_cost <= 0: continue try: @@ -159,6 +213,7 @@ async def run_ptu_reservation_rollup( __all__ = [ + "AzureCostFetcher", "PTU_ROLLUP_JOB_ID", "PTU_SENTINEL_API_KEY", "RollupResult", diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 80d5f6c4c44..656d2f5e86f 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -187,6 +187,11 @@ class UISettings(BaseModel): description="If true, enables admin-registered PTU reservations and daily flat-cost attribution on team daily spend. Governs the /ptu_reservation CRUD endpoints, the daily rollup job, the PTU Reservations UI page, and the Flat Cost column on the Usage page.", ) + enable_azure_ptu_billing_pull: bool = Field( + default=False, + description="If true and enable_ptu_cost_attribution is also true, reservations with cost_source='azure_billing' fetch daily flat cost from the Azure Cost Management API instead of the manual PTU * cost/day formula. Requires general_settings.azure_ptu_billing.subscription_id and Entra ID env vars (AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET).", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -212,6 +217,7 @@ ALLOWED_UI_SETTINGS_FIELDS = { "disable_key_generate_for_org_admin", "enable_chat_ui", "enable_ptu_cost_attribution", + "enable_azure_ptu_billing_pull", } # Flags that must be synced from the persisted UISettings into @@ -226,6 +232,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS = [ "allow_vector_stores_for_team_admins", "disable_key_generate_for_org_admin", "enable_ptu_cost_attribution", + "enable_azure_ptu_billing_pull", ] # Extension point: packages outside OSS (e.g. litellm_enterprise) can diff --git a/scripts/ptu_reservation_backfill.py b/scripts/ptu_reservation_backfill.py index e73c54b8b25..8732645ed38 100644 --- a/scripts/ptu_reservation_backfill.py +++ b/scripts/ptu_reservation_backfill.py @@ -21,6 +21,29 @@ def _parse_date(s: str) -> date: return datetime.strptime(s, "%Y-%m-%d").date() +def _build_backfill_azure_fetcher(): + """Return an Azure Cost Management client when subscription_id + Entra creds are set. + + The proxy runtime builds this from general_settings; the CLI runs outside the + proxy process so it reads AZURE_SUBSCRIPTION_ID from the environment directly. + """ + subscription_id = os.environ.get("AZURE_SUBSCRIPTION_ID") + if not subscription_id: + return None + try: + from litellm.integrations.azure_cost_management import AzureCostManagementClient + from litellm.integrations.azure_cost_management.azure_cost_management_client import ( + AzureCostManagementConfig, + ) + + config = AzureCostManagementConfig.from_env(subscription_id=subscription_id) + print(f"backfill: azure fetcher enabled for subscription {subscription_id}", file=sys.stderr) + return AzureCostManagementClient(config=config) + except Exception as exc: + print(f"backfill: azure fetcher unavailable: {exc}", file=sys.stderr) + return None + + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) group = parser.add_mutually_exclusive_group(required=True) @@ -60,10 +83,15 @@ async def _run(dates: list[date]) -> int: proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) await prisma_client.connect() + + azure_fetcher = _build_backfill_azure_fetcher() + try: total_rows = 0 for target in dates: - result = await run_ptu_reservation_rollup(prisma_client, target_date=target, force=True) + result = await run_ptu_reservation_rollup( + prisma_client, target_date=target, force=True, azure_fetcher=azure_fetcher + ) print( f"[{result.day.isoformat()}] " f"reservations={result.reservations_processed} rows_written={result.rows_written}" @@ -83,6 +111,9 @@ def main() -> int: dates = _dates_from_args(args) if "PYTHONPATH" not in os.environ: sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + import logging + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") return asyncio.run(_run(dates)) diff --git a/tests/test_litellm/integrations/azure_cost_management/__init__.py b/tests/test_litellm/integrations/azure_cost_management/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py b/tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py new file mode 100644 index 00000000000..3d4dae94567 --- /dev/null +++ b/tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py @@ -0,0 +1,145 @@ +from datetime import date +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.integrations.azure_cost_management.azure_cost_management_client import ( + AzureCostManagementClient, + AzureCostManagementConfig, + AzureCostManagementError, +) + + +def _config() -> AzureCostManagementConfig: + return AzureCostManagementConfig( + subscription_id="00000000-0000-0000-0000-000000000000", + tenant_id="tenant", + client_id="client", + client_secret="secret", + ) + + +def _http_ok(payload: dict) -> MagicMock: + http = MagicMock() + response = MagicMock() + response.json.return_value = payload + http.post = AsyncMock(return_value=response) + return http + + +@pytest.mark.asyncio +async def test_get_daily_cost_returns_sum_of_rows_when_usd(): + payload = { + "properties": { + "columns": [{"name": "Cost"}, {"name": "UsageDate"}, {"name": "Currency"}], + "rows": [ + [120.0, 20260715, "USD"], + [30.0, 20260715, "USD"], + ], + } + } + http = _http_ok(payload) + + client = AzureCostManagementClient( + config=_config(), + http_handler=http, + token_provider=lambda: "fake-token", + ) + + result = await client.get_daily_cost("/subs/x/deploy/y", date(2026, 7, 15)) + + assert result == pytest.approx(150.0) + assert client.last_currency == "USD" + + +@pytest.mark.asyncio +async def test_get_daily_cost_returns_zero_when_no_rows(): + """Azure returns 200 with empty rows when reporting lag hasn't populated yet.""" + payload = {"properties": {"columns": [{"name": "Cost"}], "rows": []}} + http = _http_ok(payload) + client = AzureCostManagementClient( + config=_config(), + http_handler=http, + token_provider=lambda: "fake-token", + ) + + result = await client.get_daily_cost("/subs/x/deploy/y", date(2026, 7, 15)) + + assert result == 0.0 + assert client.last_currency is None + + +@pytest.mark.asyncio +async def test_get_daily_cost_records_non_usd_currency(): + payload = { + "properties": { + "columns": [{"name": "Cost"}, {"name": "Currency"}], + "rows": [[100.0, "EUR"]], + } + } + http = _http_ok(payload) + client = AzureCostManagementClient( + config=_config(), + http_handler=http, + token_provider=lambda: "fake-token", + ) + + result = await client.get_daily_cost("/subs/x/deploy/y", date(2026, 7, 15)) + + assert result == 100.0 + assert client.last_currency == "EUR" + + +@pytest.mark.asyncio +async def test_get_daily_cost_raises_on_http_error(): + http = MagicMock() + response = MagicMock() + response.status_code = 403 + response.text = "forbidden" + http.post = AsyncMock( + side_effect=httpx.HTTPStatusError("forbidden", request=MagicMock(), response=response) + ) + + client = AzureCostManagementClient( + config=_config(), + http_handler=http, + token_provider=lambda: "fake-token", + ) + + with pytest.raises(AzureCostManagementError) as exc: + await client.get_daily_cost("/subs/x/deploy/y", date(2026, 7, 15)) + assert "403" in str(exc.value) + + +@pytest.mark.asyncio +async def test_get_daily_cost_posts_expected_body_shape(): + payload = {"properties": {"columns": [{"name": "Cost"}], "rows": []}} + http = _http_ok(payload) + client = AzureCostManagementClient( + config=_config(), + http_handler=http, + token_provider=lambda: "fake-token", + ) + + await client.get_daily_cost("/subs/x/deploy/y", date(2026, 7, 15)) + + http.post.assert_awaited_once() + kwargs = http.post.await_args.kwargs + assert kwargs["url"].endswith("/providers/Microsoft.CostManagement/query") + assert kwargs["params"] == {"api-version": "2023-11-01"} + assert kwargs["headers"]["Authorization"] == "Bearer fake-token" + body = kwargs["json"] + assert body["type"] == "ActualCost" + assert body["timeframe"] == "Custom" + assert body["timePeriod"]["from"] == "2026-07-15T00:00:00Z" + assert body["dataset"]["filter"]["dimensions"]["values"] == ["/subs/x/deploy/y"] + + +def test_config_from_env_raises_when_creds_missing(monkeypatch): + for k in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.delenv(k, raising=False) + + with pytest.raises(AzureCostManagementError) as exc: + AzureCostManagementConfig.from_env(subscription_id="sub-x") + assert "AZURE_TENANT_ID" in str(exc.value) diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py index f20f7cbceb7..27324d91ea6 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py @@ -24,6 +24,7 @@ class _Reservation: cost_per_ptu: float | None effective_from: datetime effective_to: datetime | None + azure_resource_id: str | None = None def _r( @@ -36,6 +37,7 @@ def _r( cost_per_ptu: float | None = 200.0, effective_from: datetime | None = None, effective_to: datetime | None = None, + azure_resource_id: str | None = None, ) -> _Reservation: return _Reservation( id=id, @@ -46,6 +48,7 @@ def _r( cost_per_ptu=cost_per_ptu, effective_from=effective_from or datetime(2026, 7, 1, tzinfo=timezone.utc), effective_to=effective_to, + azure_resource_id=azure_resource_id, ) @@ -74,35 +77,97 @@ def test_days_in_month_covers_calendar_variants(): assert _days_in_month(date(2026, 7, 31)) == 31 -def test_compute_flat_cost_calendar_month_31(): +@pytest.mark.asyncio +async def test_compute_flat_cost_calendar_month_31(): r = _r(ptu_count=1, cost_per_ptu=200.0) - assert _compute_daily_flat_cost(r, date(2026, 7, 15)) == pytest.approx(200.0 / 31) + assert await _compute_daily_flat_cost(r, date(2026, 7, 15)) == pytest.approx(200.0 / 31) -def test_compute_flat_cost_calendar_month_28(): +@pytest.mark.asyncio +async def test_compute_flat_cost_calendar_month_28(): r = _r(ptu_count=1, cost_per_ptu=200.0) - assert _compute_daily_flat_cost(r, date(2026, 2, 10)) == pytest.approx(200.0 / 28) + assert await _compute_daily_flat_cost(r, date(2026, 2, 10)) == pytest.approx(200.0 / 28) -def test_compute_flat_cost_leap_february(): +@pytest.mark.asyncio +async def test_compute_flat_cost_leap_february(): r = _r(ptu_count=1, cost_per_ptu=200.0) - assert _compute_daily_flat_cost(r, date(2024, 2, 10)) == pytest.approx(200.0 / 29) + assert await _compute_daily_flat_cost(r, date(2024, 2, 10)) == pytest.approx(200.0 / 29) -def test_compute_flat_cost_scales_with_ptu_count(): - small = _compute_daily_flat_cost(_r(ptu_count=1, cost_per_ptu=200.0), date(2026, 7, 1)) - big = _compute_daily_flat_cost(_r(ptu_count=100, cost_per_ptu=200.0), date(2026, 7, 1)) +@pytest.mark.asyncio +async def test_compute_flat_cost_scales_with_ptu_count(): + small = await _compute_daily_flat_cost(_r(ptu_count=1, cost_per_ptu=200.0), date(2026, 7, 1)) + big = await _compute_daily_flat_cost(_r(ptu_count=100, cost_per_ptu=200.0), date(2026, 7, 1)) assert big == pytest.approx(small * 100) -def test_compute_flat_cost_zero_for_non_manual_source(): - r = _r(cost_source="azure_billing", ptu_count=None, cost_per_ptu=None) - assert _compute_daily_flat_cost(r, date(2026, 7, 1)) == 0.0 +@pytest.mark.asyncio +async def test_compute_flat_cost_zero_when_azure_billing_and_no_fetcher(): + r = _r(cost_source="azure_billing", ptu_count=None, cost_per_ptu=None, azure_resource_id="/x") + assert await _compute_daily_flat_cost(r, date(2026, 7, 1)) == 0.0 -def test_compute_flat_cost_zero_when_manual_fields_missing(): +@pytest.mark.asyncio +async def test_compute_flat_cost_zero_when_manual_fields_missing(): r = _r(ptu_count=None, cost_per_ptu=None) - assert _compute_daily_flat_cost(r, date(2026, 7, 1)) == 0.0 + assert await _compute_daily_flat_cost(r, date(2026, 7, 1)) == 0.0 + + +@pytest.mark.asyncio +async def test_compute_flat_cost_azure_billing_uses_fetcher_result(): + r = _r(cost_source="azure_billing", ptu_count=None, cost_per_ptu=None, azure_resource_id="/subs/x/deploy/y") + fetcher = MagicMock() + fetcher.get_daily_cost = AsyncMock(return_value=42.5) + + result = await _compute_daily_flat_cost(r, date(2026, 7, 15), azure_fetcher=fetcher) + + assert result == 42.5 + fetcher.get_daily_cost.assert_awaited_once_with("/subs/x/deploy/y", date(2026, 7, 15)) + + +@pytest.mark.asyncio +async def test_compute_flat_cost_azure_billing_fetcher_error_returns_zero(): + r = _r(cost_source="azure_billing", ptu_count=None, cost_per_ptu=None, azure_resource_id="/subs/x/deploy/y") + fetcher = MagicMock() + fetcher.get_daily_cost = AsyncMock(side_effect=RuntimeError("azure boom")) + + result = await _compute_daily_flat_cost(r, date(2026, 7, 15), azure_fetcher=fetcher) + + assert result == 0.0 + + +@pytest.mark.asyncio +async def test_compute_flat_cost_unknown_source_returns_zero(): + r = _r(cost_source="mystery") + assert await _compute_daily_flat_cost(r, date(2026, 7, 1)) == 0.0 + + +@pytest.mark.asyncio +async def test_rollup_azure_billing_reservation_writes_fetched_amount(mock_prisma): + prisma, mock_daily, mock_reservation = mock_prisma + reservation = _r( + id="res_azure", + team_id="team_x", + model="gpt-4", + cost_source="azure_billing", + ptu_count=None, + cost_per_ptu=None, + azure_resource_id="/subs/x/deploy/y", + ) + mock_reservation.find_many = AsyncMock(return_value=[reservation]) + fetcher = MagicMock() + fetcher.get_daily_cost = AsyncMock(return_value=150.0) + + result = await run_ptu_reservation_rollup( + prisma, target_date=date(2026, 7, 12), azure_fetcher=fetcher + ) + + assert result.rows_written == 1 + fetcher.get_daily_cost.assert_awaited_once_with("/subs/x/deploy/y", date(2026, 7, 12)) + create = mock_daily.upsert.await_args.kwargs["data"]["create"] + assert create["ptu_flat_cost"] == 150.0 + assert create["ptu_reservation_id"] == "res_azure" @pytest.mark.asyncio @@ -170,11 +235,12 @@ async def test_rollup_writes_expected_row(mock_prisma): @pytest.mark.asyncio -async def test_rollup_skips_azure_billing_reservations(mock_prisma): +async def test_rollup_skips_azure_billing_reservations_when_pull_disabled(mock_prisma): + """Regression: azure_billing rows must not corrupt sentinel table when pull flag off.""" prisma, mock_daily, mock_reservation = mock_prisma mock_reservation.find_many = AsyncMock( return_value=[ - _r(cost_source="azure_billing", ptu_count=None, cost_per_ptu=None), + _r(cost_source="azure_billing", ptu_count=None, cost_per_ptu=None, azure_resource_id="/x"), ] )