fix(test): update prometheus and bias_hallucination_estimator tests to match refactored API

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Sameer Kankute 2026-06-24 17:33:37 +05:30
parent 53b2724e99
commit 723b6c55b4
No known key found for this signature in database
2 changed files with 43 additions and 28 deletions

View file

@ -19,6 +19,7 @@ from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callba
from litellm.types.integrations.prometheus import (
PrometheusMetricLabels,
UserAPIKeyLabelNames,
UserAPIKeyLabelValues,
)
from litellm.types.router import ModelInfo
@ -76,12 +77,14 @@ def test_increment_deployment_cooled_down_emits_model_group(monkeypatch):
)
logger.increment_deployment_cooled_down(
litellm_model_name="gpt-4o-mini",
model_id="model-123",
api_base="https://api.openai.com",
api_provider="openai",
enum_values=UserAPIKeyLabelValues(
litellm_model_name="gpt-4o-mini",
model_id="model-123",
api_base="https://api.openai.com",
api_provider="openai",
model_group="gpt-group",
),
exception_status="429",
model_group="gpt-group",
)
labels = logger.litellm_deployment_cooled_down.labels.call_args.kwargs
@ -100,8 +103,6 @@ def test_set_litellm_deployment_state_emits_model_group(monkeypatch):
lambda metric_name: PrometheusMetricLabels.get_labels(metric_name)
)
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
logger.set_litellm_deployment_state(
state=2,
enum_values=UserAPIKeyLabelValues(
@ -158,8 +159,6 @@ def test_deployment_metrics_omit_model_group_when_flag_disabled(monkeypatch):
lambda metric_name: PrometheusMetricLabels.get_labels(metric_name)
)
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
logger.set_litellm_deployment_state(
state=0,
enum_values=UserAPIKeyLabelValues(
@ -209,11 +208,15 @@ async def test_router_cooldown_callback_separates_alias_and_underlying_model(
# api_base resolved from the underlying model, not the alias (P2).
assert gab.call_args.kwargs["model"] == "gpt-4o-mini"
for mock in (
logger.set_deployment_complete_outage,
logger.increment_deployment_cooled_down,
):
kwargs = mock.call_args.kwargs
assert kwargs["litellm_model_name"] == "gpt-4o-mini"
assert kwargs["model_group"] == "my-gpt-group"
assert kwargs["model_id"] == "test-model-id"
# set_deployment_complete_outage uses flat kwargs
outage_kwargs = logger.set_deployment_complete_outage.call_args.kwargs
assert outage_kwargs["litellm_model_name"] == "gpt-4o-mini"
assert outage_kwargs["model_group"] == "my-gpt-group"
assert outage_kwargs["model_id"] == "test-model-id"
# increment_deployment_cooled_down uses enum_values dataclass
cooled_kwargs = logger.increment_deployment_cooled_down.call_args.kwargs
ev = cooled_kwargs["enum_values"]
assert ev.litellm_model_name == "gpt-4o-mini"
assert ev.model_group == "my-gpt-group"
assert ev.model_id == "test-model-id"

View file

@ -8,6 +8,12 @@ import pytest
from litellm.exceptions import GuardrailRaisedException
from litellm.proxy.guardrails.guardrail_hooks.bias_hallucination_estimator.bias_hallucination_estimator import (
BiasHallucinationEstimatorGuardrail,
GuardrailBehaviorConfig,
GuardrailConfig,
)
from litellm.proxy.guardrails.guardrail_hooks.bias_hallucination_estimator.risk_scorer import (
RiskThresholds,
RiskWeights,
)
from litellm.proxy.guardrails.guardrail_hooks.bias_hallucination_estimator.data_sources import (
ContextDocumentDataSource,
@ -223,8 +229,7 @@ def test_risk_scorer_blocks_when_hallucination_threshold_is_crossed() -> None:
def test_risk_scorer_flags_medium_weighted_risk() -> None:
risk = RiskScorer(
bias_threshold=0.9,
hallucination_threshold=0.9,
thresholds=RiskThresholds(bias_threshold=0.9, hallucination_threshold=0.9),
).compute_risk(
bias_analysis=BiasAnalysis(
bias_detected=True,
@ -249,7 +254,7 @@ def test_risk_scorer_passes_low_risk() -> None:
def test_risk_scorer_blocks_on_high_bias_score_alone() -> None:
risk = RiskScorer(bias_threshold=0.4).compute_risk(
risk = RiskScorer(thresholds=RiskThresholds(bias_threshold=0.4)).compute_risk(
bias_analysis=BiasAnalysis(
bias_detected=True, score=0.5, patterns_found=["overconfidence"]
),
@ -280,12 +285,14 @@ def test_risk_scorer_detected_issues_prefix_by_type() -> None:
def test_risk_scorer_custom_weights_change_overall_percentage() -> None:
bias_only = RiskScorer(bias_weight=1.0, hallucination_weight=0.0).compute_risk(
bias_only = RiskScorer(
weights=RiskWeights(bias_weight=1.0, hallucination_weight=0.0)
).compute_risk(
bias_analysis=BiasAnalysis(score=0.5),
hallucination_analysis=HallucinationAnalysis(score=0.0),
)
hallucination_only = RiskScorer(
bias_weight=0.0, hallucination_weight=1.0
weights=RiskWeights(bias_weight=0.0, hallucination_weight=1.0)
).compute_risk(
bias_analysis=BiasAnalysis(score=0.0),
hallucination_analysis=HallucinationAnalysis(score=0.5),
@ -299,7 +306,9 @@ def test_risk_scorer_custom_weights_change_overall_percentage() -> None:
def test_risk_scorer_zero_weight_total_returns_zero_percentage() -> None:
risk = RiskScorer(bias_weight=0.0, hallucination_weight=0.0).compute_risk(
risk = RiskScorer(
weights=RiskWeights(bias_weight=0.0, hallucination_weight=0.0)
).compute_risk(
bias_analysis=BiasAnalysis(score=0.1),
hallucination_analysis=HallucinationAnalysis(score=0.1),
)
@ -342,7 +351,7 @@ async def test_guardrail_blocks_high_risk_response_and_logs_metadata() -> None:
async def test_guardrail_log_payload_excludes_text_snippets() -> None:
guardrail = BiasHallucinationEstimatorGuardrail(
guardrail_name="bias-hallucination",
log_only=True,
config=GuardrailConfig(behavior=GuardrailBehaviorConfig(log_only=True)),
)
request_data: dict[str, Any] = {}
@ -378,7 +387,7 @@ async def test_guardrail_log_only_flags_without_blocking() -> None:
guardrail = BiasHallucinationEstimatorGuardrail(
guardrail_name="bias-hallucination",
event_hook=GuardrailEventHooks.post_call,
log_only=True,
config=GuardrailConfig(behavior=GuardrailBehaviorConfig(log_only=True)),
)
request_data: dict[str, Any] = {}
@ -441,7 +450,9 @@ async def test_guardrail_passes_low_risk_text() -> None:
async def test_guardrail_respects_custom_violation_message() -> None:
guardrail = BiasHallucinationEstimatorGuardrail(
guardrail_name="bias-hallucination",
violation_message="Custom block message.",
config=GuardrailConfig(
behavior=GuardrailBehaviorConfig(violation_message="Custom block message.")
),
)
with pytest.raises(GuardrailRaisedException) as exc:
@ -503,9 +514,10 @@ async def test_guardrail_empty_inputs_returns_unchanged() -> None:
async def test_guardrail_check_request_enabled_detects_bias() -> None:
guardrail = BiasHallucinationEstimatorGuardrail(
guardrail_name="bias-hallucination",
check_request=True,
check_response=False,
event_hook=GuardrailEventHooks.pre_call,
config=GuardrailConfig(
behavior=GuardrailBehaviorConfig(check_request=True, check_response=False)
),
)
with pytest.raises(GuardrailRaisedException):