From db04c2fd1b59f2bad908704d4156dcc3f7b2bf70 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 29 Sep 2026 20:45:11 +0000 Subject: [PATCH] test(integration): cover key and team redaction settings on the failure path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../observability/test_failure_redaction.py | 96 +++++++++++++++++++ .../test_failure_redaction_chaos.py | 28 +++++- 2 files changed, 121 insertions(+), 3 deletions(-) diff --git a/tests/integration/observability/test_failure_redaction.py b/tests/integration/observability/test_failure_redaction.py index d14078fe261..2b745d8836d 100644 --- a/tests/integration/observability/test_failure_redaction.py +++ b/tests/integration/observability/test_failure_redaction.py @@ -682,6 +682,67 @@ def test_b10_global_on_team_permission_opt_out_keeps_raw(rig: Rig) -> None: _assert_failure_raw(rig, model, secret, _call_id(response)) +def _logging_callback_vars(flag: bool) -> dict[str, JsonValue]: + return {"logging": [{"callback_name": "generic_api", "callback_vars": {"turn_off_message_logging": flag}}]} + + +def _event_for_call(rig: Rig, model: str, call_id: str) -> dict[str, JsonValue]: + return eventually( + lambda: tuple(event for event in rig.failure_events(model) if event.get("litellm_call_id") == call_id), + lambda values: len(values) == 1, + seconds=90, + )[0] + + +def _assert_callback_vars_decision(rig: Rig, model: str, flag: bool, key: str) -> None: + ok_secret: Final = f"ok {_secret_prompt()}" + succeeded: Final = rig.proxy.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": ok_secret}]}, key=key + ) + assert succeeded.status_code == 200, succeeded.text + success_event: Final = _event_for_call(rig, model, _call_id(succeeded)) + assert (ok_secret not in json.dumps(success_event)) == flag, json.dumps(success_event) + fail_secret: Final = _secret_prompt() + failed: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": fail_secret}]}, + key=key, + ) + assert failed.status_code == 400, failed.text + if flag: + _assert_failure_redacted(rig, model, fail_secret, _call_id(failed)) + else: + _assert_failure_raw(rig, model, fail_secret, _call_id(failed)) + + +@pytest.mark.timeout(280) +@pytest.mark.parametrize("global_flag", ["on", "off"]) +@pytest.mark.parametrize("flag", [True, False], ids=["vars_true", "vars_false"]) +def test_b6_key_logging_callback_vars_drive_the_failure_decision( + request: pytest.FixtureRequest, global_flag: str, flag: bool +) -> None: + rig: Final = request.getfixturevalue("rig" if global_flag == "on" else "rig_off") + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model], metadata=_logging_callback_vars(flag)) + _assert_callback_vars_decision(rig, model, flag, key) + + +@pytest.mark.timeout(280) +@pytest.mark.parametrize("global_flag", ["on", "off"]) +@pytest.mark.parametrize("flag", [True, False], ids=["vars_true", "vars_false"]) +def test_b7_team_logging_callback_vars_drive_the_failure_decision( + request: pytest.FixtureRequest, global_flag: str, flag: bool +) -> None: + rig: Final = request.getfixturevalue("rig" if global_flag == "on" else "rig_off") + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key") + team: Final = scenario.team(metadata=_logging_callback_vars(flag)) + key: Final = scenario.key(models=[model], team_id=team) + _assert_callback_vars_decision(rig, model, flag, key) + + # --- C. Callback registration modes (YAML global on) ---------------------------------- @@ -991,6 +1052,41 @@ def test_f3_team_opt_out_permission_flip_takes_effect(rig: Rig) -> None: assert converged[0] +@pytest.mark.timeout(280) +def test_f3_key_logging_callback_vars_flip_takes_effect(rig: Rig) -> None: + secret: Final = _secret_prompt() + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model], metadata=_logging_callback_vars(True)) + denied: Final = rig.proxy.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": secret}]}, key=key + ) + assert denied.status_code == 400, denied.text + _assert_failure_redacted(rig, model, secret, _call_id(denied)) + rig.proxy.post("/key/update", {"key": key, "metadata": _logging_callback_vars(False)}) + + def now_raw() -> tuple[bool, ...]: + probe: Final = _secret_prompt() + fired: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": probe}]}, + key=key, + ) + assert fired.status_code == 400, fired.text + events: Final = eventually( + lambda: tuple( + event for event in rig.failure_events(model) if event.get("litellm_call_id") == _call_id(fired) + ), + lambda values: len(values) == 1, + seconds=20, + ) + return (probe in json.dumps(events[0]),) + + converged: Final = eventually(now_raw, lambda values: values[0], seconds=120) + assert converged[0] + + def test_f4_five_identical_failures_log_once_each(rig: Rig) -> None: secret: Final = _secret_prompt() with rig.proxy.scenario() as scenario: diff --git a/tests/integration/observability/test_failure_redaction_chaos.py b/tests/integration/observability/test_failure_redaction_chaos.py index e4c48880267..fdd4056acf9 100644 --- a/tests/integration/observability/test_failure_redaction_chaos.py +++ b/tests/integration/observability/test_failure_redaction_chaos.py @@ -86,6 +86,14 @@ def _bodies(model: str, marker: str) -> tuple[tuple[str, dict[str, JsonValue]], return chat + messages + responses +def _json_body_ok(body: bytes) -> bool: + try: + json.loads(body) + except json.JSONDecodeError: + return False + return True + + def _fire(rig: Rig, bodies: tuple[tuple[str, dict[str, JsonValue]], ...]) -> tuple[tuple[int, str | None], ...]: def call(item: tuple[str, dict[str, JsonValue]]) -> tuple[int, str | None]: try: @@ -231,16 +239,30 @@ def test_g3_proxy_restart_mid_burst(tmp_path: Path) -> None: [], ) + probes: list[tuple[int, str]] = [] # mutable-ok: readiness polls are real served requests + def served() -> tuple[int, ...]: - probe: Final = rig_two.proxy.request("POST", *bodies[BURST // 2]) + try: + probe: Final = rig_two.proxy.request("POST", *bodies[BURST // 2]) + except httpx.HTTPError: + return (-1,) + probes.append((probe.status_code, probe.headers.get("x-litellm-call-id") or "")) return (probe.status_code,) eventually(served, lambda statuses: statuses[0] == 400, seconds=60) second: Final = _fire(rig_two, bodies[BURST // 2 :]) - answered: Final = first + second + answered: Final = first + second + tuple(probes) assert all(status in (400, 500) for status, _ in answered), answered assert any(status == 400 for status, _ in answered), answered - events: Final = tuple(object_value(event) for batch in endpoint.drain() for event in json.loads(batch.body)) + batches: Final = endpoint.drain() + empty: Final = tuple(batch for batch in batches if not _json_body_ok(batch.body)) + assert all(not batch.body.strip() for batch in empty), empty + events: Final = tuple( + object_value(event) + for batch in batches + if _json_body_ok(batch.body) + for event in json.loads(batch.body) + ) call_ids: Final = tuple(str(event.get("litellm_call_id")) for event in events if event) assert len(call_ids) == len(set(call_ids)), ("duplicate events after restart", call_ids) for event in events: