litellm-mirror/tests/litellm/integrations/test_prometheus.py
Ishaan Jaff 94a553dbb2
[Feat] Emit Key, Team Budget metrics on a cron job schedule (#9528)
* _initialize_remaining_budget_metrics

* initialize_budget_metrics_cron_job

* initialize_budget_metrics_cron_job

* initialize_budget_metrics_cron_job

* test_initialize_budget_metrics_cron_job

* LITELLM_PROXY_ADMIN_NAME

* fix code qa checks

* test_initialize_budget_metrics_cron_job

* test_initialize_budget_metrics_cron_job

* pod lock manager allow dynamic cron job ID

* fix pod lock manager

* require cronjobid for PodLockManager

* fix DB_SPEND_UPDATE_JOB_NAME acquire / release lock

* add comment on prometheus logger

* add debug statements for emitting key, team budget metrics

* test_pod_lock_manager.py

* test_initialize_budget_metrics_cron_job

* initialize_budget_metrics_cron_job

* initialize_remaining_budget_metrics

* remove outdated test
2025-04-10 16:59:14 -07:00

44 lines
1.2 KiB
Python

"""
Mock prometheus unit tests, these don't rely on LLM API calls
"""
import json
import os
import sys
import pytest
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from apscheduler.schedulers.asyncio import AsyncIOScheduler
import litellm
from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES
from litellm.integrations.prometheus import PrometheusLogger
def test_initialize_budget_metrics_cron_job():
# Create a scheduler
scheduler = AsyncIOScheduler()
# Create and register a PrometheusLogger
prometheus_logger = PrometheusLogger()
litellm.callbacks = [prometheus_logger]
# Initialize the cron job
PrometheusLogger.initialize_budget_metrics_cron_job(scheduler)
# Verify that a job was added to the scheduler
jobs = scheduler.get_jobs()
assert len(jobs) == 1
# Verify job properties
job = jobs[0]
assert (
job.trigger.interval.total_seconds() / 60
== PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES
)
assert job.func.__name__ == "initialize_remaining_budget_metrics"