From d10a5a69af7ee01394a666d2ceb0af8b4530dfa9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 29 Jul 2025 17:33:24 -0700 Subject: [PATCH] fix test_mlflow_request_tags_functionality --- .../test_litellm/integrations/test_mlflow.py | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py index 6e00fb60d0..79a5fd3b79 100644 --- a/tests/test_litellm/integrations/test_mlflow.py +++ b/tests/test_litellm/integrations/test_mlflow.py @@ -9,7 +9,6 @@ sys.path.insert(0, os.path.abspath("../..")) import pytest import litellm -from litellm.integrations.mlflow import MlflowLogger @pytest.mark.asyncio @@ -23,9 +22,27 @@ async def test_mlflow_request_tags_functionality(): mock_span.request_id = "test_trace_id" mock_client.start_trace.return_value = mock_span - with patch('mlflow.tracking.MlflowClient', return_value=mock_client), \ - patch('mlflow.get_current_active_span', return_value=None): - + # Mock all MLflow-related imports to avoid requiring MLflow as a dependency + mock_mlflow_tracking = MagicMock() + mock_mlflow_tracking.MlflowClient = MagicMock(return_value=mock_client) + + mock_mlflow_entities = MagicMock() + mock_mlflow_entities.SpanStatusCode.OK = "OK" + mock_mlflow_entities.SpanStatusCode.ERROR = "ERROR" + mock_mlflow_entities.SpanType.LLM = "LLM" + + mock_mlflow = MagicMock() + mock_mlflow.get_current_active_span.return_value = None + + with patch.dict('sys.modules', { + 'mlflow': mock_mlflow, + 'mlflow.tracking': mock_mlflow_tracking, + 'mlflow.entities': mock_mlflow_entities, + 'mlflow.tracing.utils': MagicMock(), + }): + # Now we can safely import MlflowLogger + from litellm.integrations.mlflow import MlflowLogger + # Create MlflowLogger instance mlflow_logger = MlflowLogger() litellm.callbacks = [mlflow_logger] @@ -56,5 +73,3 @@ async def test_mlflow_request_tags_functionality(): assert tags_param == expected_tags, f"Expected tags {expected_tags}, got {tags_param}" print("✅ Request tags properly transformed and passed to MLflow trace") - -