cleanup

tests/loggers/test_mlflow_logger.py

@@ -193,26 +193,22 @@ 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')
 
         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')