From ce679eed9b039a09cfbc8ad2fd953c29d84755c9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 1 Apr 2026 08:50:07 -0700 Subject: [PATCH] fix: update test_increment_remaining_budget_metrics to pass budget reset times in metadata instead of patching get_team_object/get_key_object --- .../test_prometheus_logging_callbacks.py | 134 +++++++++--------- 1 file changed, 64 insertions(+), 70 deletions(-) diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 967ef2a5fec..ecaca644d3a 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -282,85 +282,79 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): # Create a future budget reset time for testing future_reset_time_team = datetime.now() + timedelta(hours=10) future_reset_time_key = datetime.now() + timedelta(hours=12) - # Mock the get_team_object and get_key_object functions to return objects with budget reset times - with patch( - "litellm.proxy.auth.auth_checks.get_team_object" - ) as mock_get_team, patch( - "litellm.proxy.auth.auth_checks.get_key_object" - ) as mock_get_key: - mock_get_team.return_value = MagicMock(budget_reset_at=future_reset_time_team) - mock_get_key.return_value = MagicMock(budget_reset_at=future_reset_time_key) - litellm_params = { - "metadata": { - "user_api_key_team_spend": 50, - "user_api_key_team_max_budget": 100, - "user_api_key_spend": 25, - "user_api_key_max_budget": 75, - } + litellm_params = { + "metadata": { + "user_api_key_team_spend": 50, + "user_api_key_team_max_budget": 100, + "user_api_key_spend": 25, + "user_api_key_max_budget": 75, + "user_api_key_team_budget_reset_at": future_reset_time_team.isoformat(), + "user_api_key_budget_reset_at": future_reset_time_key.isoformat(), } + } - await prometheus_logger._increment_remaining_budget_metrics( - user_api_team="team1", - user_api_team_alias="team_alias1", - user_api_key="key1", - user_api_key_alias="alias1", - litellm_params=litellm_params, - response_cost=10, - ) + await prometheus_logger._increment_remaining_budget_metrics( + user_api_team="team1", + user_api_team_alias="team_alias1", + user_api_key="key1", + user_api_key_alias="alias1", + litellm_params=litellm_params, + response_cost=10, + ) - # Test remaining budget metrics - prometheus_logger.litellm_remaining_team_budget_metric.labels.assert_called_once_with( - team="team1", team_alias="team_alias1" - ) - prometheus_logger.litellm_remaining_team_budget_metric.labels().set.assert_called_once_with( - 40 # 100 - (50 + 10) - ) + # Test remaining budget metrics + prometheus_logger.litellm_remaining_team_budget_metric.labels.assert_called_once_with( + team="team1", team_alias="team_alias1" + ) + prometheus_logger.litellm_remaining_team_budget_metric.labels().set.assert_called_once_with( + 40 # 100 - (50 + 10) + ) - prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once_with( - hashed_api_key="key1", api_key_alias="alias1" - ) - prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.assert_called_once_with( - 40 # 75 - (25 + 10) - ) + prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once_with( + hashed_api_key="key1", api_key_alias="alias1" + ) + prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.assert_called_once_with( + 40 # 75 - (25 + 10) + ) - # Test max budget metrics - prometheus_logger.litellm_team_max_budget_metric.labels.assert_called_once_with( - team="team1", team_alias="team_alias1" - ) - prometheus_logger.litellm_team_max_budget_metric.labels().set.assert_called_once_with( - 100 - ) + # Test max budget metrics + prometheus_logger.litellm_team_max_budget_metric.labels.assert_called_once_with( + team="team1", team_alias="team_alias1" + ) + prometheus_logger.litellm_team_max_budget_metric.labels().set.assert_called_once_with( + 100 + ) - prometheus_logger.litellm_api_key_max_budget_metric.labels.assert_called_once_with( - hashed_api_key="key1", api_key_alias="alias1" - ) - prometheus_logger.litellm_api_key_max_budget_metric.labels().set.assert_called_once_with( - 75 - ) + prometheus_logger.litellm_api_key_max_budget_metric.labels.assert_called_once_with( + hashed_api_key="key1", api_key_alias="alias1" + ) + prometheus_logger.litellm_api_key_max_budget_metric.labels().set.assert_called_once_with( + 75 + ) - # Test remaining hours metrics - prometheus_logger.litellm_team_budget_remaining_hours_metric.labels.assert_called_once_with( - team="team1", team_alias="team_alias1" - ) - # The remaining hours should be approximately 10 (with some small difference due to test execution time) - remaining_hours_call = prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.call_args[ - 0 - ][ - 0 - ] - assert 9.9 <= remaining_hours_call <= 10.0 + # Test remaining hours metrics + prometheus_logger.litellm_team_budget_remaining_hours_metric.labels.assert_called_once_with( + team="team1", team_alias="team_alias1" + ) + # The remaining hours should be approximately 10 (with some small difference due to test execution time) + remaining_hours_call = prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.call_args[ + 0 + ][ + 0 + ] + assert 9.9 <= remaining_hours_call <= 10.0 - prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels.assert_called_once_with( - hashed_api_key="key1", api_key_alias="alias1" - ) - # The remaining hours should be approximately 10 (with some small difference due to test execution time) - remaining_hours_call = prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels().set.call_args[ - 0 - ][ - 0 - ] - assert 11.9 <= remaining_hours_call <= 12.0 + prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels.assert_called_once_with( + hashed_api_key="key1", api_key_alias="alias1" + ) + # The remaining hours should be approximately 12 (with some small difference due to test execution time) + remaining_hours_call = prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels().set.call_args[ + 0 + ][ + 0 + ] + assert 11.9 <= remaining_hours_call <= 12.0 def test_set_latency_metrics(prometheus_logger):