Repository: Lightning-AI/pytorch-lightning
File: tests/loggers/test_mlflow.py (Python)

Code before:

def test_mlflow_logger_dirs_creation(tmpdir):
    """ Test that the logger creates the folders and files in the right place. """
    if not importlib.util.find_spec('mlflow'):
        pytest.xfail("test for explicit file creation requires mlflow dependency to be installed.")

    assert not os.listdir(tmpdir)
    logger = MLFlowLogger('test', save_dir=tmpdir)
    assert logger.save_dir == tmpdir
    assert set(os.listdir(tmpdir)) == {'.trash'}
    run_id = logger.run_id
    exp_id = logger.experiment_id

    # multiple experiment calls should not lead to new experiment folders
    for i in range(2):
        _ = logger.experiment
        assert set(os.listdir(tmpdir)) == {'.trash', exp_id}
        assert set(os.listdir(tmpdir / exp_id)) == {run_id, 'meta.yaml'}

    model = EvalModelTemplate()
    trainer = Trainer(default_root_dir=tmpdir, logger=logger, max_epochs=1, limit_val_batches=3,
                      log_gpu_memory=True)
    trainer.fit(model)
    assert set(os.listdir(tmpdir / exp_id)) == {run_id, 'meta.yaml'}
    assert 'epoch' in os.listdir(tmpdir / exp_id / run_id / 'metrics')
    assert set(os.listdir(tmpdir / exp_id / run_id / 'params')) == model.hparams.keys()
    assert trainer.checkpoint_callback.dirpath == (tmpdir / exp_id / run_id / 'checkpoints')
    assert set(os.listdir(trainer.checkpoint_callback.dirpath)) == {'epoch=0-step=9.ckpt'}


@mock.patch('pytorch_lightning.loggers.mlflow.mlflow')
@mock.patch('pytorch_lightning.loggers.mlflow.MlflowClient')
def test_mlflow_experiment_id_retrieved_once(client, mlflow, tmpdir):
    """
    Test that the logger experiment_id retrieved only once.
    """
    logger = MLFlowLogger('test', save_dir=tmpdir)
    _ = logger.experiment
    _ = logger.experiment
    _ = logger.experiment
    assert logger.experiment.get_experiment_by_name.call_count == 1


@mock.patch('pytorch_lightning.loggers.mlflow.mlflow')
@mock.patch('pytorch_lightning.loggers.mlflow.MlflowClient')
def test_mlflow_logger_with_unexpected_characters(client, mlflow, tmpdir):
    """
    Test that the logger raises warning with special characters not accepted by MLFlow.
    """
    logger = MLFlowLogger('test', save_dir=tmpdir)
    metrics = {'[some_metric]': 10}

Review comment (performance):
```suggestion
        limit_train_batches=3,
        limit_val_batches=3,
```

Code after:

def test_mlflow_logger_dirs_creation(tmpdir):
    """ Test that the logger creates the folders and files in the right place. """
    if not _MLFLOW_AVAILABLE:
        pytest.xfail("test for explicit file creation requires mlflow dependency to be installed.")

    assert not os.listdir(tmpdir)
    logger = MLFlowLogger('test', save_dir=tmpdir)
    assert logger.save_dir == tmpdir
    assert set(os.listdir(tmpdir)) == {'.trash'}
    run_id = logger.run_id
    exp_id = logger.experiment_id

    # multiple experiment calls should not lead to new experiment folders
    for i in range(2):
        _ = logger.experiment
        assert set(os.listdir(tmpdir)) == {'.trash', exp_id}
        assert set(os.listdir(tmpdir / exp_id)) == {run_id, 'meta.yaml'}

    model = EvalModelTemplate()
    trainer = Trainer(
        default_root_dir=tmpdir,
        logger=logger,
        max_epochs=1,
        limit_val_batches=3,
        log_gpu_memory=True,
    )
    trainer.fit(model)
    assert set(os.listdir(tmpdir / exp_id)) == {run_id, 'meta.yaml'}
    assert 'epoch' in os.listdir(tmpdir / exp_id / run_id / 'metrics')
    assert set(os.listdir(tmpdir / exp_id / run_id / 'params')) == model.hparams.keys()
    assert trainer.checkpoint_callback.dirpath == (tmpdir / exp_id / run_id / 'checkpoints')
    assert set(os.listdir(trainer.checkpoint_callback.dirpath)) == {'epoch=0-step=9.ckpt'}


@mock.patch('pytorch_lightning.loggers.mlflow.mlflow')
@mock.patch('pytorch_lightning.loggers.mlflow.MlflowClient')
def test_mlflow_experiment_id_retrieved_once(client, mlflow, tmpdir):
    """
    Test that the logger experiment_id retrieved only once.
    """
    logger = MLFlowLogger('test', save_dir=tmpdir)
    _ = logger.experiment
    _ = logger.experiment
    _ = logger.experiment
    assert logger.experiment.get_experiment_by_name.call_count == 1


@mock.patch('pytorch_lightning.loggers.mlflow.mlflow')
@mock.patch('pytorch_lightning.loggers.mlflow.MlflowClient')
def test_mlflow_logger_with_unexpected_characters(client, mlflow, tmpdir):