test(prometheus): add unit tests for org budget metrics

tests/test_litellm/integrations/test_prometheus_user_team_metrics.py

@@ -2,7 +2,7 @@
 Unit tests for Prometheus user and team count metrics
 """
 from datetime import datetime, timezone
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from prometheus_client import REGISTRY
@@ -523,3 +523,150 @@ async def test_set_user_budget_metrics_after_api_request_inf_when_genuinely_no_b
     assert actual_value == float("inf"), (
         "remaining_user_budget_metric should be +Inf when user truly has no budget"
     )
+
+
+# ---------------------------------------------------------------------------
+# Org budget metric tests
+# ---------------------------------------------------------------------------
+
+
+def test_org_budget_metrics_initialized(prometheus_logger):
+    """Test that the 3 org budget gauge metrics are initialized."""
+    assert hasattr(prometheus_logger, "litellm_remaining_org_budget_metric")
+    assert hasattr(prometheus_logger, "litellm_org_max_budget_metric")
+    assert hasattr(prometheus_logger, "litellm_org_budget_remaining_hours_metric")
+    assert prometheus_logger.litellm_remaining_org_budget_metric is not None
+    assert prometheus_logger.litellm_org_max_budget_metric is not None
+    assert prometheus_logger.litellm_org_budget_remaining_hours_metric is not None
+
+
+def test_set_org_budget_metrics_remaining_budget(prometheus_logger):
+    """_set_org_budget_metrics sets remaining budget gauge correctly."""
+    prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_max_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock()
+
+    prometheus_logger._set_org_budget_metrics(
+        org_id="org-abc",
+        org_alias="my-org",
+        spend=200.0,
+        max_budget=500.0,
+        budget_reset_at=None,
+    )
+
+    set_call = prometheus_logger.litellm_remaining_org_budget_metric.labels().set
+    set_call.assert_called_once()
+    actual = set_call.call_args[0][0]
+    assert abs(actual - 300.0) < 0.01, f"Expected 300.0, got {actual}"
+
+
+def test_set_org_budget_metrics_max_budget(prometheus_logger):
+    """_set_org_budget_metrics sets max budget gauge when max_budget is not None."""
+    prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_max_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock()
+
+    prometheus_logger._set_org_budget_metrics(
+        org_id="org-abc",
+        org_alias="my-org",
+        spend=100.0,
+        max_budget=1000.0,
+        budget_reset_at=None,
+    )
+
+    prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with(
+        1000.0
+    )
+
+
+def test_set_org_budget_metrics_no_max_budget(prometheus_logger):
+    """_set_org_budget_metrics does not set max budget gauge when max_budget is None."""
+    prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_max_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock()
+
+    prometheus_logger._set_org_budget_metrics(
+        org_id="org-abc",
+        org_alias="my-org",
+        spend=50.0,
+        max_budget=None,
+        budget_reset_at=None,
+    )
+
+    prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_not_called()
+
+
+def test_set_org_budget_metrics_remaining_hours(prometheus_logger):
+    """_set_org_budget_metrics sets remaining hours gauge when budget_reset_at is set."""
+    prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_max_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock()
+
+    future_reset = datetime(2099, 1, 1, tzinfo=timezone.utc)
+    prometheus_logger._set_org_budget_metrics(
+        org_id="org-abc",
+        org_alias="my-org",
+        spend=10.0,
+        max_budget=500.0,
+        budget_reset_at=future_reset,
+    )
+
+    prometheus_logger.litellm_org_budget_remaining_hours_metric.labels().set.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_set_org_budget_metrics_after_api_request(prometheus_logger):
+    """_set_org_budget_metrics_after_api_request fetches from DB and sets gauges."""
+    import sys
+
+    prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_max_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock()
+
+    budget_mock = MagicMock()
+    budget_mock.max_budget = 1000.0
+    budget_mock.budget_reset_at = datetime(2099, 1, 1, tzinfo=timezone.utc)
+
+    org_mock = MagicMock()
+    org_mock.organization_id = "org-xyz"
+    org_mock.organization_alias = "test-org"
+    org_mock.spend = 300.0
+    org_mock.litellm_budget_table = budget_mock
+    org_mock.model_dump.return_value = {}
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock(
+        return_value=org_mock
+    )
+
+    mock_proxy_server = MagicMock()
+    mock_proxy_server.prisma_client = mock_prisma
+
+    with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
+        await prometheus_logger._set_org_budget_metrics_after_api_request(
+            org_id="org-xyz",
+            response_cost=0.0,
+        )
+
+    prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once()
+    prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with(
+        1000.0
+    )
+    prometheus_logger.litellm_org_budget_remaining_hours_metric.labels().set.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_set_org_budget_metrics_after_api_request_no_org_id(prometheus_logger):
+    """_set_org_budget_metrics_after_api_request is a no-op when org_id is None."""
+    prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_max_budget_metric = MagicMock()
+    prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock()
+
+    await prometheus_logger._set_org_budget_metrics_after_api_request(
+        org_id=None,
+        response_cost=1.0,
+    )
+
+    prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_not_called()
+    prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_not_called()
+    prometheus_logger.litellm_org_budget_remaining_hours_metric.labels().set.assert_not_called()