diff --git a/tests/loggers/test_mlflow_logger.py b/tests/loggers/test_mlflow_logger.py index 573ba789e4..4f784c479a 100644 --- a/tests/loggers/test_mlflow_logger.py +++ b/tests/loggers/test_mlflow_logger.py @@ -193,18 +193,15 @@ def test_mlflow_experiment_init_existing_composer_run(monkeypatch): def test_mlflow_logger_uses_env_var_run_name(monkeypatch): """Test that MLFlowLogger uses the 'RUN_NAME' environment variable if set.""" mlflow = pytest.importorskip('mlflow') - MlflowClient = mlflow.tracking.MlflowClient # Import locally after ensuring mlflow is available - # Set up mocks monkeypatch.setattr(mlflow, 'set_tracking_uri', MagicMock()) monkeypatch.setattr(mlflow, 'start_run', MagicMock()) mock_create_run = MagicMock(return_value=MagicMock(info=MagicMock(run_id='mock-run-id'))) - # Use 'patch' to mock MlflowClient inside the test function where it's used with patch('mlflow.tracking.MlflowClient') as MockClient: MockClient.return_value.create_run = mock_create_run - from composer.loggers.mlflow_logger import MLFlowLogger # Import the logger here after the patch + from composer.loggers.mlflow_logger import MLFlowLogger mock_state = MagicMock() mock_state.run_name = 'dummy-run-name' monkeypatch.setenv('RUN_NAME', 'env-run-name') @@ -212,7 +209,6 @@ def test_mlflow_logger_uses_env_var_run_name(monkeypatch): test_logger = MLFlowLogger() test_logger.init(state=mock_state, logger=MagicMock()) - # Assertion to check the correct environment variable is used assert test_logger.tags['run_name'] == 'env-run-name', "Logger should use the run name from the environment variable." monkeypatch.delenv('RUN_NAME')