def test_mlflow_log_model(mlflow_mock, log_model, tmp_path):
    """Test that the logger creates the folders and files in the right place."""
    client = mlflow_mock.tracking.MlflowClient

    # Get model, logger, trainer and train
    model = BoringModel()
    logger = MLFlowLogger("test", save_dir=str(tmp_path), log_model=log_model)
    logger = mock_mlflow_run_creation(logger, experiment_id="test-id")

    trainer = Trainer(
        default_root_dir=tmp_path,
        logger=logger,
        max_epochs=2,
        limit_train_batches=3,
        limit_val_batches=3,
    )
    trainer.fit(model)

    if log_model == "all":
        # Checkpoint log
        assert client.return_value.log_artifact.call_count == 2
        # Metadata and aliases log
        assert client.return_value.log_artifacts.call_count == 2

    elif log_model is True:
        # Checkpoint log
        client.return_value.log_artifact.assert_called_once()
        # Metadata and aliases log
        client.return_value.log_artifacts.assert_called_once()

    elif log_model is False:
        # Checkpoint log
        assert not client.return_value.log_artifact.called
        # Metadata and aliases log
        assert not client.return_value.log_artifacts.called
Test that the logger creates the folders and files in the right place.
tests/tests_pytorch/loggers/test_mlflow.py