scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(ptu): fetch azure_billing flat cost via Cost Management API

Browse source 
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.
This commit is contained in:
Yucheng Zhu 2026-07-20 16:50:23 -07:00
parent 1b85426a6a
commit 152e233722
11 changed files with 573 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/integrations/azure_cost_management/__init__.py Normal file
View file

@ -0,0 +1,6 @@
from litellm.integrations.azure_cost_management.azure_cost_management_client import (
AzureCostManagementClient,
AzureCostManagementError,
)
__all__ = ["AzureCostManagementClient", "AzureCostManagementError"]

191
litellm/integrations/azure_cost_management/azure_cost_management_client.py Normal file
View file

@ -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

2
litellm/proxy/dev_config.yaml
View file

@ -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

4
litellm/proxy/management_endpoints/ptu_reservation_endpoints.py
View file

@ -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:

41
litellm/proxy/proxy_server.py
View file

@ -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,

79
litellm/proxy/spend_tracking/ptu_reservation_rollup.py
View file

@ -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",

7
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -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

33
scripts/ptu_reservation_backfill.py
View file

@ -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))

0
tests/test_litellm/integrations/azure_cost_management/__init__.py Normal file
View file

145
tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py Normal file
View file

@ -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)

98
tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py
View file

@ -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"),
]
)