From 3db6d2a1ed35d1dce04e686288cef10078da6df3 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 6 Dec 2025 07:26:34 -0800 Subject: [PATCH] Reapply Langfuse logger test mock setup fix (#17591) Reapplies the fix from commit a885e215432d94fe2f9443c871f96c4565aed509 that was reverted in 6c9556be67e18f35454457d1dc25c70eec17013f. The original revert was done because the test was flaky and giving false negatives. This fix properly mocks the Langfuse client to ensure the test can correctly verify that _log_langfuse_v2 converts None usage values to 0. Changes: - Add mock_langfuse_client.client attribute to prevent errors during init - Add trace_id to mock_langfuse_generation for proper return value handling - Remove redundant mock setup code - Explicitly set logger.Langfuse to mock client after initialization - Set logger.langfuse_sdk_version to ensure _supports_* methods work correctly --- .../integrations/test_langfuse.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index b9009d20056..5610ac8fff6 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -35,8 +35,11 @@ class TestLangfuseUsageDetails(unittest.TestCase): # Create mock objects self.mock_langfuse_client = MagicMock() + # Mock the client attribute to prevent errors during logger initialization + self.mock_langfuse_client.client = MagicMock() self.mock_langfuse_trace = MagicMock() self.mock_langfuse_generation = MagicMock() + self.mock_langfuse_generation.trace_id = "test-trace-id" # Setup the trace and generation chain self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation @@ -63,22 +66,13 @@ class TestLangfuseUsageDetails(unittest.TestCase): sys.modules["langfuse"] = self.mock_langfuse sys.modules["langfuse"].Langfuse = self.mock_langfuse_class - # Mock the Langfuse client - self.mock_langfuse_client = MagicMock() - self.mock_langfuse_trace = MagicMock() - self.mock_langfuse_generation = MagicMock() - - # Setup the trace and generation chain - self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation - self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace - - # Mock the Langfuse class - self.mock_langfuse_class = MagicMock() - self.mock_langfuse_class.return_value = self.mock_langfuse_client - self.mock_langfuse.Langfuse = self.mock_langfuse_class - # Create the logger self.logger = LangFuseLogger() + + # Explicitly set the Langfuse client to our mock + self.logger.Langfuse = self.mock_langfuse_client + # Ensure langfuse_sdk_version is set correctly for _supports_* methods + self.logger.langfuse_sdk_version = "3.0.0" # Add the log_event_on_langfuse method to the instance def log_event_on_langfuse(