def test_mlflow_logger_experiment_calls_with_synchronous(mlflow_mock, tmp_path, synchronous):
    """Test that the logger calls methods on the mlflow experiment with the specified synchronous flag."""

    time = mlflow_mock.entities.time
    metric = mlflow_mock.entities.Metric
    param = mlflow_mock.entities.Param
    time.return_value = 1

    mlflow_client = mlflow_mock.tracking.MlflowClient.return_value
    mlflow_client.get_experiment_by_name.return_value = None
    logger = MLFlowLogger(
        "test", save_dir=str(tmp_path), artifact_location="my_artifact_location", synchronous=synchronous
    )

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

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

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

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

    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 with the specified synchronous flag.
tests/tests_pytorch/loggers/test_mlflow.py