def test_mlflow_run_id_setting(mlflow_mock, tmp_path):
    """Test that the run_id argument uses the provided run_id."""
    client = mlflow_mock.tracking.MlflowClient

    run = MagicMock()
    run.info.run_id = "run-id"
    run.info.experiment_id = "experiment-id"

    # simulate existing run
    client.return_value.get_run = MagicMock(return_value=run)

    # run_id exists uses the existing run
    logger = MLFlowLogger("test", run_id=run.info.run_id, save_dir=str(tmp_path))
    _ = logger.experiment
    client.return_value.get_run.assert_called_with(run.info.run_id)
    assert logger.experiment_id == run.info.experiment_id
    assert logger.run_id == run.info.run_id
    client.reset_mock(return_value=True)
Test that the run_id argument uses the provided run_id.
tests/tests_pytorch/loggers/test_mlflow.py