diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 010d9f863c..8da7dd8917 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -268,10 +268,10 @@ class TestLangfuseUsageDetails(unittest.TestCase): Test that _log_langfuse_v2 correctly handles None values in the usage object by converting them to 0, preventing validation errors. """ - # Reset the mock to ensure clean state - self.mock_langfuse_client.reset_mock() - self.mock_langfuse_trace.reset_mock() - self.mock_langfuse_generation.reset_mock() + # Reset the mock to ensure clean state; clear side_effect so return_value takes effect + self.mock_langfuse_client.reset_mock(side_effect=True) + self.mock_langfuse_trace.reset_mock(side_effect=True) + self.mock_langfuse_generation.reset_mock(side_effect=True) # Re-setup the trace and generation chain with clean state self.mock_langfuse_generation.trace_id = "test-trace-id" @@ -283,12 +283,14 @@ class TestLangfuseUsageDetails(unittest.TestCase): # Ensure trace returns our mock self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace self.logger.Langfuse = self.mock_langfuse_client - + with patch( "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", side_effect=lambda generation_params, **kwargs: generation_params, create=True, - ) as mock_add_prompt_params: + ) as mock_add_prompt_params, patch.object( + self.logger, "_supports_prompt", return_value=True + ): # Create a mock response object with usage information containing None values response_obj = MagicMock() response_obj.usage = MagicMock()