def test_mlflow_log_dir(mlflow_mock, tmp_path):
    """Test that the trainer saves checkpoints in the logger's save dir."""
    client = mlflow_mock.tracking.MlflowClient

    # simulate experiment creation with mlflow client mock
    run = MagicMock()
    run.info.run_id = "run-id"
    client.return_value.get_experiment_by_name = MagicMock(return_value=None)
    client.return_value.create_experiment = MagicMock(return_value="exp-id")
    client.return_value.create_run = MagicMock(return_value=run)

    # test construction of default log dir path
    logger = MLFlowLogger("test", save_dir=str(tmp_path))
    assert logger.save_dir == str(tmp_path)
    assert logger.version == "run-id"
    assert logger.name == "exp-id"

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path, logger=logger, max_epochs=1, limit_train_batches=1, limit_val_batches=3
    )
    assert trainer.log_dir == logger.save_dir
    trainer.fit(model)
    assert trainer.checkpoint_callback.dirpath == str(tmp_path / "exp-id" / "run-id" / "checkpoints")
    assert set(os.listdir(trainer.checkpoint_callback.dirpath)) == {"epoch=0-step=1.ckpt"}
    assert trainer.log_dir == logger.save_dir
Test that the trainer saves checkpoints in the logger's save dir.
tests/tests_pytorch/loggers/test_mlflow.py