def test_mlflow_logger_experiment_calls(mlflow_mock, tmp_path):
    """Test that the logger calls methods on the mlflow experiment correctly."""
    time = mlflow_mock.entities.time
    metric = mlflow_mock.entities.Metric
    param = mlflow_mock.entities.Param

    time.return_value = 1

    logger = MLFlowLogger("test", save_dir=str(tmp_path), artifact_location="my_artifact_location")
    logger._mlflow_client.get_experiment_by_name.return_value = None

    params = {"test": "test_param"}
    logger.log_hyperparams(params)

    logger.experiment.log_batch.assert_called_once_with(
        run_id=logger.run_id, params=[param(key="test", value="test_param")]
    )
    param.assert_called_with(key="test", value="test_param")

    metrics = {"some_metric": 10}
    logger.log_metrics(metrics)

    logger.experiment.log_batch.assert_called_with(
        run_id=logger.run_id, metrics=[metric(key="some_metric", value=10, timestamp=1000, step=0)]
    )
    metric.assert_called_with(key="some_metric", value=10, timestamp=1000, step=0)

    logger._mlflow_client.create_experiment.assert_called_once_with(
        name="test", artifact_location="my_artifact_location"
    )
Test that the logger calls methods on the mlflow experiment correctly.
tests/tests_pytorch/loggers/test_mlflow.py