mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
b76def0e5d
`pytest.raises(Exception)` with no `match=` passes on any error that broad. A TypeError from a refactor, a botched fixture, an import that moved: all of them read as the rejection the test claims to police, so the test goes green for the wrong reason and stays green after the behaviour it guards is gone. PT011 closes that gap for the 317 sites B017 could not reach, because B017 only fires on a single-statement body with no `as e` binding. Each pattern here is the message the code actually raised, recorded by running the sites under a plugin that logged the concrete type and text per call site, so the assertions describe observed behaviour rather than a guess. Where a site raises more than one message across its parametrize cases, the pattern is an alternation of what was seen; where the exception carries an empty `str()` and puts the text on `.message`, the site keeps a narrow `noqa` with the reason. PT014 removes four parametrize cases that were listed twice. The duplicate re-runs an assertion that already passed, and it usually marks a case someone meant to vary and forgot to edit.
361 lines
12 KiB
Python
361 lines
12 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from litellm.integrations.levo.levo import LevoConfig, LevoLogger
|
|
from litellm.integrations.opentelemetry import OpenTelemetryConfig
|
|
|
|
# Try to import OpenTelemetry packages, skip tests if not available
|
|
try:
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
|
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
|
|
InMemorySpanExporter,
|
|
)
|
|
|
|
OPENTELEMETRY_AVAILABLE = True
|
|
except ImportError:
|
|
OPENTELEMETRY_AVAILABLE = False
|
|
|
|
|
|
class TestLevoConfig(unittest.TestCase):
|
|
"""Unit tests for LevoLogger configuration."""
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
|
|
"LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
|
|
},
|
|
)
|
|
def test_get_levo_config_with_all_required_vars(self):
|
|
"""Test get_levo_config() with all required environment variables."""
|
|
config = LevoLogger.get_levo_config()
|
|
|
|
# Verify headers include all three values
|
|
self.assertIn("Authorization=Bearer test-api-key", config.otlp_auth_headers)
|
|
self.assertIn("x-levo-organization-id=test-org-id", config.otlp_auth_headers)
|
|
self.assertIn("x-levo-workspace-id=test-workspace-id", config.otlp_auth_headers)
|
|
|
|
# Verify endpoint uses provided collector URL exactly as-is
|
|
self.assertEqual(config.endpoint, "https://collector.levo.ai")
|
|
|
|
# Verify protocol is otlp_http
|
|
self.assertEqual(config.protocol, "otlp_http")
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
|
|
"LEVOAI_COLLECTOR_URL": "https://custom.collector.com",
|
|
},
|
|
)
|
|
def test_get_levo_config_with_custom_collector_url(self):
|
|
"""Test get_levo_config() with custom collector URL."""
|
|
config = LevoLogger.get_levo_config()
|
|
|
|
# Verify endpoint uses custom URL exactly as provided
|
|
self.assertEqual(config.endpoint, "https://custom.collector.com")
|
|
self.assertEqual(config.protocol, "otlp_http")
|
|
|
|
@patch.dict("os.environ", {}, clear=True)
|
|
def test_get_levo_config_missing_api_key(self):
|
|
"""Test get_levo_config() raises ValueError when LEVOAI_API_KEY is missing."""
|
|
with pytest.raises(ValueError, match="LEVOAI_API_KEY"):
|
|
LevoLogger.get_levo_config()
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
},
|
|
clear=True,
|
|
)
|
|
def test_get_levo_config_missing_org_id(self):
|
|
"""Test get_levo_config() raises ValueError when LEVOAI_ORG_ID is missing."""
|
|
with pytest.raises(ValueError, match="LEVOAI_ORG_ID"):
|
|
LevoLogger.get_levo_config()
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
},
|
|
clear=True,
|
|
)
|
|
def test_get_levo_config_missing_workspace_id(self):
|
|
"""Test get_levo_config() raises ValueError when LEVOAI_WORKSPACE_ID is missing."""
|
|
with pytest.raises(ValueError, match="LEVOAI_WORKSPACE_ID"):
|
|
LevoLogger.get_levo_config()
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
|
|
},
|
|
clear=True,
|
|
)
|
|
def test_get_levo_config_missing_collector_url(self):
|
|
"""Test get_levo_config() raises ValueError when LEVOAI_COLLECTOR_URL is missing."""
|
|
with pytest.raises(ValueError, match="LEVOAI_COLLECTOR_URL"):
|
|
LevoLogger.get_levo_config()
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
|
|
"LEVOAI_COLLECTOR_URL": "http://localhost:4318",
|
|
},
|
|
)
|
|
def test_get_levo_config_with_http_endpoint(self):
|
|
"""Test get_levo_config() with HTTP endpoint."""
|
|
config = LevoLogger.get_levo_config()
|
|
|
|
# Should use HTTP endpoint exactly as provided
|
|
self.assertEqual(config.endpoint, "http://localhost:4318")
|
|
self.assertEqual(config.protocol, "otlp_http")
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
|
|
"LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
|
|
},
|
|
)
|
|
def test_levo_config_headers_format(self):
|
|
"""Test that OTLP headers are formatted correctly."""
|
|
config = LevoLogger.get_levo_config()
|
|
|
|
# Verify headers contain all required parts
|
|
self.assertIn("Authorization=Bearer test-api-key", config.otlp_auth_headers)
|
|
self.assertIn("x-levo-organization-id=test-org-id", config.otlp_auth_headers)
|
|
self.assertIn("x-levo-workspace-id=test-workspace-id", config.otlp_auth_headers)
|
|
|
|
# Verify headers are comma-separated
|
|
header_parts = config.otlp_auth_headers.split(",")
|
|
self.assertEqual(len(header_parts), 3)
|
|
|
|
|
|
class TestLevoIntegration(unittest.TestCase):
|
|
"""Integration tests for LevoLogger."""
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
|
|
"LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
|
|
},
|
|
)
|
|
@pytest.mark.skipif(
|
|
not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed"
|
|
)
|
|
@patch(
|
|
"litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy"
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_levo_logger_health_check_healthy(self, mock_init_proxy):
|
|
"""Test health check returns healthy status when config is valid."""
|
|
# Mock the proxy initialization to avoid importing proxy code
|
|
mock_init_proxy.return_value = None
|
|
|
|
config = LevoLogger.get_levo_config()
|
|
otel_config = OpenTelemetryConfig(
|
|
exporter=config.protocol,
|
|
endpoint=config.endpoint,
|
|
headers=config.otlp_auth_headers,
|
|
)
|
|
|
|
# Create tracer provider with in-memory exporter
|
|
tracer_provider = TracerProvider()
|
|
tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
|
|
|
|
levo_logger = LevoLogger(
|
|
config=otel_config, callback_name="levo", tracer_provider=tracer_provider
|
|
)
|
|
|
|
# Run health check
|
|
result = await levo_logger.async_health_check()
|
|
|
|
self.assertEqual(result["status"], "healthy")
|
|
self.assertIn("message", result)
|
|
|
|
@patch.dict("os.environ", {}, clear=True)
|
|
def test_levo_logger_health_check_unhealthy(self):
|
|
"""Test health check returns unhealthy status when required vars are missing."""
|
|
# Try to create logger without required env vars
|
|
# This should fail during config, but we can test health check logic
|
|
with pytest.raises(ValueError, match='LEVOAI_API_KEY environment variable is required for Levo'):
|
|
LevoLogger.get_levo_config()
|
|
|
|
@patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LEVOAI_API_KEY": "test-api-key",
|
|
"LEVOAI_ORG_ID": "test-org-id",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
|
|
"LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
|
|
},
|
|
)
|
|
@pytest.mark.skipif(
|
|
not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed"
|
|
)
|
|
@patch(
|
|
"litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy"
|
|
)
|
|
def test_levo_logger_callback_name(self, mock_init_proxy):
|
|
"""Test that callback_name is properly set and used."""
|
|
# Mock the proxy initialization to avoid importing proxy code
|
|
mock_init_proxy.return_value = None
|
|
|
|
config = LevoLogger.get_levo_config()
|
|
otel_config = OpenTelemetryConfig(
|
|
exporter=config.protocol,
|
|
endpoint=config.endpoint,
|
|
headers=config.otlp_auth_headers,
|
|
)
|
|
|
|
# Create tracer provider with in-memory exporter
|
|
tracer_provider = TracerProvider()
|
|
tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
|
|
|
|
levo_logger = LevoLogger(
|
|
config=otel_config, callback_name="levo", tracer_provider=tracer_provider
|
|
)
|
|
|
|
# Verify callback_name attribute
|
|
self.assertEqual(levo_logger.callback_name, "levo")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"env_vars, expected_headers_contains, expected_endpoint, expected_protocol",
|
|
[
|
|
pytest.param(
|
|
{
|
|
"LEVOAI_API_KEY": "test-key",
|
|
"LEVOAI_ORG_ID": "test-org",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace",
|
|
"LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
|
|
},
|
|
[
|
|
"Authorization=Bearer test-key",
|
|
"x-levo-organization-id=test-org",
|
|
"x-levo-workspace-id=test-workspace",
|
|
],
|
|
"https://collector.levo.ai",
|
|
"otlp_http",
|
|
id="collector URL with all required vars",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"LEVOAI_API_KEY": "key-123",
|
|
"LEVOAI_ORG_ID": "org-456",
|
|
"LEVOAI_WORKSPACE_ID": "workspace-789",
|
|
"LEVOAI_COLLECTOR_URL": "https://custom.example.com",
|
|
},
|
|
[
|
|
"Authorization=Bearer key-123",
|
|
"x-levo-organization-id=org-456",
|
|
"x-levo-workspace-id=workspace-789",
|
|
],
|
|
"https://custom.example.com",
|
|
"otlp_http",
|
|
id="custom collector URL",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"LEVOAI_API_KEY": "key-123",
|
|
"LEVOAI_ORG_ID": "org-456",
|
|
"LEVOAI_WORKSPACE_ID": "workspace-789",
|
|
"LEVOAI_COLLECTOR_URL": "http://localhost:9999",
|
|
},
|
|
["Authorization=Bearer key-123"],
|
|
"http://localhost:9999",
|
|
"otlp_http",
|
|
id="custom HTTP endpoint",
|
|
),
|
|
],
|
|
)
|
|
def test_get_levo_config_parametrized(
|
|
monkeypatch,
|
|
env_vars,
|
|
expected_headers_contains,
|
|
expected_endpoint,
|
|
expected_protocol,
|
|
):
|
|
"""Parametrized tests for get_levo_config() with various configurations."""
|
|
# Clear all Levo-related env vars first to ensure clean state
|
|
for key in [
|
|
"LEVOAI_API_KEY",
|
|
"LEVOAI_ORG_ID",
|
|
"LEVOAI_WORKSPACE_ID",
|
|
"LEVOAI_COLLECTOR_URL",
|
|
"LEVOAI_ENV_NAME",
|
|
]:
|
|
monkeypatch.delenv(key, raising=False)
|
|
|
|
for key, value in env_vars.items():
|
|
monkeypatch.setenv(key, value)
|
|
|
|
config = LevoLogger.get_levo_config()
|
|
|
|
assert isinstance(config, LevoConfig)
|
|
assert config.endpoint == expected_endpoint
|
|
assert config.protocol == expected_protocol
|
|
|
|
# Verify all expected header parts are present
|
|
for header_part in expected_headers_contains:
|
|
assert header_part in config.otlp_auth_headers
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"missing_var",
|
|
[
|
|
pytest.param("LEVOAI_API_KEY", id="missing API key"),
|
|
pytest.param("LEVOAI_ORG_ID", id="missing org ID"),
|
|
pytest.param("LEVOAI_WORKSPACE_ID", id="missing workspace ID"),
|
|
pytest.param("LEVOAI_COLLECTOR_URL", id="missing collector URL"),
|
|
],
|
|
)
|
|
def test_get_levo_config_missing_required_vars(monkeypatch, missing_var):
|
|
"""Test that missing required environment variables raise ValueError."""
|
|
# Clear all Levo-related env vars
|
|
for key in [
|
|
"LEVOAI_API_KEY",
|
|
"LEVOAI_ORG_ID",
|
|
"LEVOAI_WORKSPACE_ID",
|
|
"LEVOAI_COLLECTOR_URL",
|
|
]:
|
|
monkeypatch.delenv(key, raising=False)
|
|
|
|
# Set all required vars except the missing one
|
|
required_vars = {
|
|
"LEVOAI_API_KEY": "test-key",
|
|
"LEVOAI_ORG_ID": "test-org",
|
|
"LEVOAI_WORKSPACE_ID": "test-workspace",
|
|
"LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
|
|
}
|
|
required_vars.pop(missing_var)
|
|
|
|
for key, value in required_vars.items():
|
|
monkeypatch.setenv(key, value)
|
|
|
|
with pytest.raises(ValueError, match=missing_var):
|
|
LevoLogger.get_levo_config()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|