chore: import upstream snapshot with attribution

tests/pyfunc/docker/test_docker.py

@@ -0,0 +1,157 @@
+import difflib
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from unittest import mock
+
+import pytest
+import sklearn
+import sklearn.neighbors
+from packaging.version import Version
+
+import mlflow
+from mlflow.environment_variables import _MLFLOW_RUN_SLOW_TESTS
+from mlflow.models import Model
+from mlflow.models.docker_utils import build_image_from_context
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.utils import PYTHON_VERSION
+from mlflow.utils.env_manager import CONDA, LOCAL, VIRTUALENV
+from mlflow.version import VERSION
+
+from tests.pyfunc.docker.conftest import RESOURCE_DIR, get_released_mlflow_version
+
+
+def _get_mlflow_install_specifier():
+    if Version(VERSION).is_devrelease:
+        return "https://github.com/mlflow/mlflow/archive/refs/heads/master.zip"
+    return f"mlflow=={VERSION}"
+
+
+def assert_dockerfiles_equal(actual_dockerfile_path: Path, expected_dockerfile_path: Path):
+    actual_dockerfile = actual_dockerfile_path.read_text()
+    expected_dockerfile = (
+        expected_dockerfile_path
+        .read_text()
+        .replace("${{ MLFLOW_INSTALL }}", _get_mlflow_install_specifier())
+        .replace("${{ PYTHON_VERSION }}", PYTHON_VERSION)
+    )
+    assert actual_dockerfile == expected_dockerfile, (
+        "Generated Dockerfile does not match expected one. Diff:\n"
+        + "\n".join(
+            difflib.unified_diff(expected_dockerfile.splitlines(), actual_dockerfile.splitlines())
+        )
+    )
+
+
+def save_model(tmp_path):
+    knn_model = sklearn.neighbors.KNeighborsClassifier()
+    model_path = os.path.join(tmp_path, "model")
+    mlflow.sklearn.save_model(
+        knn_model,
+        path=model_path,
+        pip_requirements=[
+            f"mlflow=={get_released_mlflow_version()}",
+            f"scikit-learn=={sklearn.__version__}",
+        ],  # Skip requirements inference for speed up
+    )
+    return model_path
+
+
+def add_spark_flavor_to_model(model_path):
+    model_config_path = os.path.join(model_path, "MLmodel")
+    model = Model.load(model_config_path)
+    model.add_flavor("spark", spark_version="3.5.0")
+    model.save(model_config_path)
+
+
+@dataclass
+class Param:
+    expected_dockerfile: str
+    env_manager: str | None = None
+    mlflow_home: str | None = None
+    install_mlflow: bool = False
+    # If True, image is built with --model-uri param
+    specify_model_uri: bool = True
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        Param(expected_dockerfile="Dockerfile_default"),
+        Param(expected_dockerfile="Dockerfile_default", env_manager=LOCAL),
+        Param(expected_dockerfile="Dockerfile_java_flavor", env_manager=VIRTUALENV),
+        Param(expected_dockerfile="Dockerfile_conda", env_manager=CONDA),
+        Param(install_mlflow=True, expected_dockerfile="Dockerfile_install_mlflow"),
+        Param(mlflow_home=".", expected_dockerfile="Dockerfile_with_mlflow_home"),
+        Param(specify_model_uri=False, expected_dockerfile="Dockerfile_no_model_uri"),
+    ],
+)
+def test_build_image(tmp_path, params):
+    model_uri = save_model(tmp_path) if params.specify_model_uri else None
+
+    backend = get_flavor_backend(model_uri, docker_build=True, env_manager=params.env_manager)
+
+    # Copy the context dir to a temp dir so we can verify the generated Dockerfile
+    def _build_image_with_copy(context_dir, image_name):
+        shutil.copytree(context_dir, dst_dir)
+        # Build the image if the slow-tests flag is enabled
+        if _MLFLOW_RUN_SLOW_TESTS.get():
+            for _ in range(3):
+                try:
+                    # Docker image build is unstable on GitHub Actions, retry up to 3 times
+                    build_image_from_context(context_dir, image_name)
+                    break
+                except RuntimeError:
+                    pass
+            else:
+                raise RuntimeError("Docker image build failed.")
+
+    dst_dir = tmp_path / "context"
+    with mock.patch(
+        "mlflow.models.docker_utils.build_image_from_context",
+        side_effect=_build_image_with_copy,
+    ):
+        backend.build_image(
+            model_uri=model_uri,
+            image_name="test_image",
+            mlflow_home=params.mlflow_home,
+            install_mlflow=params.install_mlflow,
+        )
+
+    actual = dst_dir / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / params.expected_dockerfile
+    assert_dockerfiles_equal(actual, expected)
+
+
+def test_generate_dockerfile_for_java_flavor(tmp_path):
+    model_path = save_model(tmp_path)
+    add_spark_flavor_to_model(model_path)
+
+    backend = get_flavor_backend(model_path, docker_build=True, env_manager=None)
+
+    backend.generate_dockerfile(
+        model_uri=model_path,
+        output_dir=tmp_path,
+    )
+
+    actual = tmp_path / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / "Dockerfile_java_flavor"
+    assert_dockerfiles_equal(actual, expected)
+
+
+def test_generate_dockerfile_for_custom_image(tmp_path):
+    model_path = save_model(tmp_path)
+    add_spark_flavor_to_model(model_path)
+
+    backend = get_flavor_backend(model_path, docker_build=True, env_manager=None)
+
+    backend.generate_dockerfile(
+        base_image="quay.io/jupyter/scipy-notebook:latest",
+        model_uri=model_path,
+        output_dir=tmp_path,
+    )
+
+    actual = tmp_path / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / "Dockerfile_custom_scipy"
+    assert_dockerfiles_equal(actual, expected)