Merge pull request #21419 from BerriAI/fix/langfuse-test-supports-prompt-flakiness

fix(test): prevent flaky failure in test_log_langfuse_v2_handles_null_usage_values

Author: jquinter

tests/test_litellm/integrations/test_langfuse.py

@@ -268,10 +268,10 @@ def test_log_langfuse_v2_handles_null_usage_values(self):
         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 @@ def test_log_langfuse_v2_handles_null_usage_values(self):
         # 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()