Skip to content

test(audit-logs): let monkeypatch own the audit log and s3 callback globals #37842

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuneng-berri merged 1 commit into from
Aug 22, 2026
+41 −44
Merged

test(audit-logs): let monkeypatch own the audit log and s3 callback globals #37842

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion test-quality-budget.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
"limit": 768
},
"TQ005": {
"limit": 2832
"limit": 2810
},
"TQ006": {
"limit": 34
Expand Down
83 changes: 40 additions & 43 deletions tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,9 @@


@pytest.fixture(autouse=True)
def reset_audit_log_callbacks():
"""Reset audit_log_callbacks before and after each test."""
original = litellm.audit_log_callbacks
litellm.audit_log_callbacks = []
yield
litellm.audit_log_callbacks = original
def reset_audit_log_callbacks(monkeypatch: pytest.MonkeyPatch) -> None:
"""Every test starts with no audit log callbacks registered."""
monkeypatch.setattr(litellm, "audit_log_callbacks", [])


def _make_audit_log(
Expand Down Expand Up @@ -115,10 +112,10 @@ def test_handles_none_values(self):

class TestDispatchAuditLogToCallbacks:
@pytest.mark.asyncio
async def test_dispatches_to_custom_logger_instance(self):
async def test_dispatches_to_custom_logger_instance(self, monkeypatch: pytest.MonkeyPatch):
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock()
litellm.audit_log_callbacks = [mock_logger]
monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger])

audit_log = _make_audit_log()
await _dispatch_audit_log_to_callbacks(audit_log)
Expand All @@ -132,18 +129,18 @@ async def test_dispatches_to_custom_logger_instance(self):
assert payload["action"] == "created"

@pytest.mark.asyncio
async def test_no_dispatch_when_callbacks_empty(self):
litellm.audit_log_callbacks = []
async def test_no_dispatch_when_callbacks_empty(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "audit_log_callbacks", [])
audit_log = _make_audit_log()
# Should return immediately without error
await _dispatch_audit_log_to_callbacks(audit_log)

@pytest.mark.asyncio
async def test_resolves_string_callback(self):
async def test_resolves_string_callback(self, monkeypatch: pytest.MonkeyPatch):
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock()

litellm.audit_log_callbacks = ["s3_v2"]
monkeypatch.setattr(litellm, "audit_log_callbacks", ["s3_v2"])

with patch(
"litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback",
Expand All @@ -156,22 +153,22 @@ async def test_resolves_string_callback(self):
mock_logger.async_log_audit_log_event.assert_called_once()

@pytest.mark.asyncio
async def test_nonblocking_on_callback_failure(self):
async def test_nonblocking_on_callback_failure(self, monkeypatch: pytest.MonkeyPatch):
"""Callback errors should not propagate."""
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock(
side_effect=RuntimeError("boom")
)
litellm.audit_log_callbacks = [mock_logger]
monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger])

audit_log = _make_audit_log()
# Should not raise
await _dispatch_audit_log_to_callbacks(audit_log)
await asyncio.sleep(0.1)

@pytest.mark.asyncio
async def test_skips_unresolvable_string_callback(self):
litellm.audit_log_callbacks = ["nonexistent_callback"]
async def test_skips_unresolvable_string_callback(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "audit_log_callbacks", ["nonexistent_callback"])

with patch(
"litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback",
Expand All @@ -184,10 +181,10 @@ async def test_skips_unresolvable_string_callback(self):

class TestCreateAuditLogForUpdateWithCallbacks:
@pytest.mark.asyncio
async def test_dispatches_to_callbacks_after_db_write(self):
async def test_dispatches_to_callbacks_after_db_write(self, monkeypatch: pytest.MonkeyPatch):
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock()
litellm.audit_log_callbacks = [mock_logger]
monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger])

with (
patch("litellm.proxy.proxy_server.premium_user", True),
Expand All @@ -206,10 +203,10 @@ async def test_dispatches_to_callbacks_after_db_write(self):
mock_logger.async_log_audit_log_event.assert_called_once()

@pytest.mark.asyncio
async def test_no_dispatch_when_not_premium(self):
async def test_no_dispatch_when_not_premium(self, monkeypatch: pytest.MonkeyPatch):
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock()
litellm.audit_log_callbacks = [mock_logger]
monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger])

with (
patch("litellm.proxy.proxy_server.premium_user", False),
Expand All @@ -224,10 +221,10 @@ async def test_no_dispatch_when_not_premium(self):
mock_prisma.db.litellm_auditlog.create.assert_not_called()

@pytest.mark.asyncio
async def test_no_dispatch_when_store_audit_logs_false(self):
async def test_no_dispatch_when_store_audit_logs_false(self, monkeypatch: pytest.MonkeyPatch):
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock()
litellm.audit_log_callbacks = [mock_logger]
monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger])

with patch("litellm.store_audit_logs", False):
audit_log = _make_audit_log()
Expand All @@ -237,11 +234,11 @@ async def test_no_dispatch_when_store_audit_logs_false(self):
mock_logger.async_log_audit_log_event.assert_not_called()

@pytest.mark.asyncio
async def test_dispatches_even_when_prisma_client_is_none(self):
async def test_dispatches_even_when_prisma_client_is_none(self, monkeypatch: pytest.MonkeyPatch):
"""Callbacks should fire even if DB is unavailable."""
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock()
litellm.audit_log_callbacks = [mock_logger]
monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger])

with (
patch("litellm.proxy.proxy_server.premium_user", True),
Expand All @@ -256,11 +253,11 @@ async def test_dispatches_even_when_prisma_client_is_none(self):
mock_logger.async_log_audit_log_event.assert_called_once()

@pytest.mark.asyncio
async def test_dispatches_even_when_db_write_fails(self):
async def test_dispatches_even_when_db_write_fails(self, monkeypatch: pytest.MonkeyPatch):
"""Callbacks should fire even if the DB write raises."""
mock_logger = MagicMock(spec=CustomLogger)
mock_logger.async_log_audit_log_event = AsyncMock()
litellm.audit_log_callbacks = [mock_logger]
monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger])

with (
patch("litellm.proxy.proxy_server.premium_user", True),
Expand Down Expand Up @@ -384,21 +381,21 @@ class TestS3AuditCallbackParamsDecoupling:
S3Logger instance, distinct from the singleton serving normal logs."""

@pytest.fixture(autouse=True)
def _isolate_caches_and_globals(self):
def _isolate_caches_and_globals(self, monkeypatch: pytest.MonkeyPatch):
from litellm.litellm_core_utils import litellm_logging as ll_logging
from litellm.proxy.management_helpers import audit_logs as ll_audit_logs

original_s3 = litellm.s3_callback_params
original_audit = getattr(litellm, "s3_audit_callback_params", None)
monkeypatch.setattr(litellm, "s3_callback_params", litellm.s3_callback_params)
monkeypatch.setattr(
litellm, "s3_audit_callback_params", getattr(litellm, "s3_audit_callback_params", None)
)
ll_audit_logs._audit_log_callback_cache.clear()
ll_logging._in_memory_loggers.clear()
yield
litellm.s3_callback_params = original_s3
litellm.s3_audit_callback_params = original_audit
ll_audit_logs._audit_log_callback_cache.clear()
ll_logging._in_memory_loggers.clear()

def test_opt_in_constructs_separate_instance_with_audit_config(self):
def test_opt_in_constructs_separate_instance_with_audit_config(self, monkeypatch: pytest.MonkeyPatch):
"""Audit config set → audit resolver returns a fresh S3Logger pointing
at the audit bucket, distinct from the normal-log singleton."""
from litellm.integrations.s3_v2 import S3Logger
Expand All @@ -409,8 +406,8 @@ def test_opt_in_constructs_separate_instance_with_audit_config(self):
_resolve_audit_log_callback,
)

litellm.s3_callback_params = {"s3_bucket_name": "normal-bucket"}
litellm.s3_audit_callback_params = {"s3_bucket_name": "audit-bucket"}
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"})
monkeypatch.setattr(litellm, "s3_audit_callback_params", {"s3_bucket_name": "audit-bucket"})

with patch("asyncio.create_task"):
audit_instance = _resolve_audit_log_callback("s3_v2")
Expand All @@ -426,7 +423,7 @@ def test_opt_in_constructs_separate_instance_with_audit_config(self):
assert audit_instance.s3_bucket_name == "audit-bucket"
assert normal_instance.s3_bucket_name == "normal-bucket"

def test_opt_out_preserves_singleton_behavior(self):
def test_opt_out_preserves_singleton_behavior(self, monkeypatch: pytest.MonkeyPatch):
"""No `s3_audit_callback_params` → audit and normal share the singleton
(existing behavior, regression guard)."""
from litellm.integrations.s3_v2 import S3Logger
Expand All @@ -437,8 +434,8 @@ def test_opt_out_preserves_singleton_behavior(self):
_resolve_audit_log_callback,
)

litellm.s3_callback_params = {"s3_bucket_name": "shared-bucket"}
litellm.s3_audit_callback_params = None
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "shared-bucket"})
monkeypatch.setattr(litellm, "s3_audit_callback_params", None)

with patch("asyncio.create_task"):
normal_instance = _init_custom_logger_compatible_class(
Expand All @@ -452,7 +449,7 @@ def test_opt_out_preserves_singleton_behavior(self):
assert id(audit_instance) == id(normal_instance)
assert audit_instance.s3_bucket_name == "shared-bucket"

def test_empty_dict_opts_in(self):
def test_empty_dict_opts_in(self, monkeypatch: pytest.MonkeyPatch):
"""`s3_audit_callback_params = {}` is opt-in (truthy-by-presence) and
produces a separate instance with no bucket configured (env/IAM-only)."""
from litellm.integrations.s3_v2 import S3Logger
Expand All @@ -463,8 +460,8 @@ def test_empty_dict_opts_in(self):
_resolve_audit_log_callback,
)

litellm.s3_callback_params = {"s3_bucket_name": "normal-bucket"}
litellm.s3_audit_callback_params = {}
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"})
monkeypatch.setattr(litellm, "s3_audit_callback_params", {})

with patch("asyncio.create_task"):
audit_instance = _resolve_audit_log_callback("s3_v2")
Expand All @@ -478,7 +475,7 @@ def test_empty_dict_opts_in(self):
assert audit_instance.s3_bucket_name is None
assert normal_instance.s3_bucket_name == "normal-bucket"

def test_reset_audit_log_callback_cache_clears_audit_instance(self):
def test_reset_audit_log_callback_cache_clears_audit_instance(self, monkeypatch: pytest.MonkeyPatch):
"""`reset_audit_log_callback_cache()` must drop the cached audit
instance so a config reload picks up the new params."""
from litellm.proxy.management_helpers.audit_logs import (
Expand All @@ -487,15 +484,15 @@ def test_reset_audit_log_callback_cache_clears_audit_instance(self):
reset_audit_log_callback_cache,
)

litellm.s3_audit_callback_params = {"s3_bucket_name": "first"}
monkeypatch.setattr(litellm, "s3_audit_callback_params", {"s3_bucket_name": "first"})
with patch("asyncio.create_task"):
first = _resolve_audit_log_callback("s3_v2")
assert first is not None and "s3_v2" in _audit_log_callback_cache

reset_audit_log_callback_cache()
assert "s3_v2" not in _audit_log_callback_cache

litellm.s3_audit_callback_params = {"s3_bucket_name": "second"}
monkeypatch.setattr(litellm, "s3_audit_callback_params", {"s3_bucket_name": "second"})
second = _resolve_audit_log_callback("s3_v2")
assert second is not None
assert id(second) != id(first)
Expand Down
Loading