How to use the mlflow.utils.file_utils.TempDir function in mlflow

def test_path_params():
    with TempDir() as tmp:
        dest_path = tmp.path()
        data_file = "s3://path.test/resources/data_file.csv"
        defaults = {
            "constants": {"type": "uri", "default": "s3://path.test/b1"},
            "data": {"type": "path", "default": data_file}
        }
        entry_point = EntryPoint("entry_point_name", defaults, "command_name script.py")

        with mock.patch("mlflow.data.download_uri") as download_uri_mock:
            final_1, extra_1 = entry_point.compute_parameters({}, None)
            assert (final_1 == {"constants": "s3://path.test/b1", "data": data_file})
            assert (extra_1 == {})
            assert download_uri_mock.call_count == 0

        with mock.patch("mlflow.data.download_uri") as download_uri_mock:
            user_2 = {"alpha": 0.001, "constants": "s3://path.test/b_two"}

def test_model_log(xgb_model, model_path):
    old_uri = mlflow.get_tracking_uri()
    model = xgb_model.model
    with TempDir(chdr=True, remove_on_exit=True) as tmp:
        for should_start_run in [False, True]:
            try:
                mlflow.set_tracking_uri("test")
                if should_start_run:
                    mlflow.start_run()

                artifact_path = "model"
                conda_env = os.path.join(tmp.path(), "conda_env.yaml")
                _mlflow_conda_env(conda_env, additional_pip_deps=["xgboost"])

                mlflow.xgboost.log_model(
                        xgb_model=model,
                        artifact_path=artifact_path,
                        conda_env=conda_env)
                model_uri = "runs:/{run_id}/{artifact_path}".format(
                    run_id=mlflow.active_run().info.run_id,

def test_build_image_with_relative_model_path_calls_expected_azure_routines(
        sklearn_model):
    with TempDir(chdr=True):
        model_path = "model"
        mlflow.sklearn.save_model(sk_model=sklearn_model, path=model_path)
        with AzureMLMocks() as aml_mocks:
            workspace = get_azure_workspace()
            mlflow.azureml.build_image(model_uri=model_path, workspace=workspace)

            assert aml_mocks["register_model"].call_count == 1
            assert aml_mocks["create_image"].call_count == 1

def test_spark_module_model_save_with_relative_path_and_valid_sample_input_produces_mleap_flavor(
        spark_model_iris):
    with TempDir(chdr=True) as tmp:
        model_path = os.path.basename(tmp.path("model"))
        mlflow_model = Model()
        sparkm.save_model(spark_model=spark_model_iris.model,
                          path=model_path,
                          sample_input=spark_model_iris.spark_df,
                          mlflow_model=mlflow_model)
        assert mleap.FLAVOR_NAME in mlflow_model.flavors

        config_path = os.path.join(model_path, "MLmodel")
        assert os.path.exists(config_path)
        config = Model.load(config_path)
        assert mleap.FLAVOR_NAME in config.flavors

def test_cli_build_image_with_relative_model_path_calls_expected_azure_routines(sklearn_model):
    with TempDir(chdr=True):
        model_path = "model"
        mlflow.sklearn.save_model(sk_model=sklearn_model, path=model_path)

        with AzureMLMocks() as aml_mocks:
            result = CliRunner(env={"LC_ALL": "en_US.UTF-8", "LANG": "en_US.UTF-8"}).invoke(
                    mlflow.azureml.cli.commands,
                    [
                        'build-image',
                        '-m', model_path,
                        '-w', 'test_workspace',
                        '-i', 'image_name',
                        '-n', 'model_name',
                    ])
            assert result.exit_code == 0

            assert aml_mocks["register_model"].call_count == 1

def test_gbt():
    old_uri = tracking.get_tracking_uri()
    with TempDir(chdr=False, remove_on_exit=True) as tmp:
        try:
            diamonds = tmp.path("diamonds")
            artifacts = tmp.path("artifacts")
            os.mkdir(diamonds)
            os.mkdir(artifacts)
            tracking.set_tracking_uri(artifacts)
            mlflow.set_experiment("test-experiment")
            # Download the diamonds dataset via mlflow run
            run(".", entry_point="main", version=None,
                parameters={"dest-dir": diamonds},
                mode="local", cluster_spec=None, git_username=None, git_password=None,
                use_conda=True, storage_dir=None)

            # Run the main gbt app via mlflow
            submitted_run = run(
                "apps/gbt-regression", entry_point="main", version=None,

def test_download_artifacts(local_artifact_repo, dst_path):
    artifact_rel_path = "test.txt"
    artifact_text = "hello world!"
    empty_dir_path = "empty_dir"
    with TempDir(chdr=True) as local_dir:
        if dst_path:
            os.mkdir(dst_path)
        artifact_src_path = local_dir.path(artifact_rel_path)
        os.mkdir(local_dir.path(empty_dir_path))
        with open(artifact_src_path, "w") as f:
            f.write(artifact_text)
        local_artifact_repo.log_artifacts(local_dir.path())
        result = local_artifact_repo.download_artifacts(artifact_path=artifact_rel_path,
                                                        dst_path=dst_path)
        assert open(result).read() == artifact_text
        result = local_artifact_repo.download_artifacts(artifact_path="", dst_path=dst_path)
        empty_dir_dst_path = os.path.join(result, empty_dir_path)
        assert os.path.isdir(empty_dir_dst_path)
        assert len(os.listdir(empty_dir_dst_path)) == 0

def test_build_image_includes_mlflow_home_as_file_dependency_if_specified(
        sklearn_model, model_path):
    def mock_create_dockerfile(output_path, *args, **kwargs):
        # pylint: disable=unused-argument
        with open(output_path, "w") as f:
            f.write("Dockerfile contents")

    mlflow.sklearn.save_model(sk_model=sklearn_model, path=model_path)
    with AzureMLMocks() as aml_mocks, TempDir() as tmp,\
            mock.patch("mlflow.azureml._create_dockerfile") as create_dockerfile_mock:
        create_dockerfile_mock.side_effect = mock_create_dockerfile

        # Write a mock `setup.py` file to the mlflow home path so that it will be recognized
        # as a viable MLflow source directory during the image build process
        mlflow_home = tmp.path()
        with open(os.path.join(mlflow_home, "setup.py"), "w") as f:
            f.write("setup instructions")

        workspace = get_azure_workspace()
        mlflow.azureml.build_image(
            model_uri=model_path, workspace=workspace, mlflow_home=mlflow_home)

        assert len(create_dockerfile_mock.call_args_list) == 1
        _, create_dockerfile_kwargs = create_dockerfile_mock.call_args_list[0]
        # The path to MLflow that is referenced by the Docker container may differ from the

def test_entry_point_compute_command():
    """
    Tests that EntryPoint correctly computes the command to execute in order to run the entry point.
    """
    project = load_project()
    entry_point = project.get_entry_point("greeter")
    with TempDir() as tmp:
        storage_dir = tmp.path()
        command = entry_point.compute_command({"name": "friend", "excitement": 10}, storage_dir)
        assert command == "python greeter.py hi friend --excitement 10"
        with pytest.raises(ExecutionException):
            entry_point.compute_command({}, storage_dir)
        # Test shell escaping
        name_value = "friend; echo 'hi'"
        command = entry_point.compute_command({"name": name_value}, storage_dir)
        assert command == "python greeter.py %s %s" % (shlex_quote("hi"), shlex_quote(name_value))

def _get_flavor_backend(model_uri, **kwargs):
    with TempDir() as tmp:
        if ModelsArtifactRepository.is_models_uri(model_uri):
            underlying_model_uri = ModelsArtifactRepository.get_underlying_uri(model_uri)
        else:
            underlying_model_uri = model_uri
        local_path = _download_artifact_from_uri(posixpath.join(underlying_model_uri, "MLmodel"),
                                                 output_path=tmp.path())
        model = Model.load(local_path)
    flavor_name, flavor_backend = get_flavor_backend(model, **kwargs)
    if flavor_backend is None:
        raise Exception("No suitable flavor backend was found for the model.")
    _logger.info("Selected backend for flavor '%s'", flavor_name)
    return flavor_backend