chore: import upstream snapshot with attribution

tests/utils/test_model_utils.py

@@ -0,0 +1,128 @@
+import os
+import sys
+from unittest import mock
+
+import pytest
+import sklearn.linear_model as logreg_module
+from sklearn import datasets
+
+import mlflow.sklearn
+import mlflow.utils.model_utils as mlflow_model_utils
+from mlflow.environment_variables import MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING
+from mlflow.exceptions import MlflowException
+from mlflow.models import Model
+from mlflow.utils.file_utils import TempDir
+from mlflow.utils.model_utils import env_var_tracker
+
+
+@pytest.fixture(scope="module")
+def sklearn_knn_model():
+    iris = datasets.load_iris()
+    X = iris.data[:, :2]  # we only take the first two features.
+    y = iris.target
+    logreg_model = logreg_module.LogisticRegression()
+    logreg_model.fit(X, y)
+    return logreg_model
+
+
+@pytest.fixture
+def model_path(tmp_path):
+    return os.path.join(tmp_path, "model")
+
+
+def test_get_flavor_configuration_throws_exception_when_requested_flavor_is_missing(
+    model_path, sklearn_knn_model
+):
+    mlflow.sklearn.save_model(sk_model=sklearn_knn_model, path=model_path)
+
+    # The saved model contains the "sklearn" flavor, so this call should succeed
+    sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
+        model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME
+    )
+    assert sklearn_flavor_config is not None
+
+
+def test_get_flavor_configuration_with_present_flavor_returns_expected_configuration(
+    sklearn_knn_model, model_path
+):
+    mlflow.sklearn.save_model(sk_model=sklearn_knn_model, path=model_path)
+
+    sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
+        model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME
+    )
+    model_config = Model.load(os.path.join(model_path, "MLmodel"))
+    assert sklearn_flavor_config == model_config.flavors[mlflow.sklearn.FLAVOR_NAME]
+
+
+def test_add_code_to_system_path(sklearn_knn_model, model_path):
+    mlflow.sklearn.save_model(
+        sk_model=sklearn_knn_model,
+        path=model_path,
+        code_paths=[
+            "tests/utils/test_resources/dummy_module.py",
+            "tests/utils/test_resources/dummy_package",
+        ],
+    )
+
+    sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
+        model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME
+    )
+    with TempDir(chdr=True):
+        # Load the model from a new directory that is not a parent of the source code path to
+        # verify that source code paths and their subdirectories are resolved correctly
+        with pytest.raises(ModuleNotFoundError, match="No module named 'dummy_module'"):
+            import dummy_module
+
+        mlflow_model_utils._add_code_from_conf_to_system_path(model_path, sklearn_flavor_config)
+        import dummy_module  # noqa: F401
+
+    # If this raises an exception it's because dummy_package.test imported
+    # dummy_package.operator and not the built-in operator module. This only
+    # happens if MLflow is misconfiguring the system path.
+    from dummy_package import base  # noqa: F401
+
+    # Ensure that the custom tests/utils/test_resources/dummy_package/pandas.py is not
+    # overwriting the 3rd party `pandas` package
+    assert "dummy_package" in sys.modules
+    assert "pandas" in sys.modules
+    assert "site-packages" in sys.modules["pandas"].__file__
+
+
+def test_add_code_to_system_path_not_copyable_file(sklearn_knn_model, model_path):
+    with mock.patch("builtins.open", side_effect=OSError("[Errno 95] Operation not supported")):
+        with pytest.raises(MlflowException, match=r"Failed to copy the specified code path"):
+            mlflow.sklearn.save_model(
+                sk_model=sklearn_knn_model,
+                path=model_path,
+                code_paths=["tests/utils/test_resources/dummy_module.py"],
+            )
+
+
+def test_env_var_tracker(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_HOST", "host")
+    assert "DATABRICKS_HOST" in os.environ
+    assert "TEST_API_KEY" not in os.environ
+
+    with env_var_tracker() as tracked_env_names:
+        assert os.environ["DATABRICKS_HOST"] == "host"
+        monkeypatch.setenv("TEST_API_KEY", "key")
+        # accessed env var is tracked
+        assert os.environ.get("TEST_API_KEY") == "key"
+        # test non-existing env vars fetched by `get` are not tracked
+        os.environ.get("INVALID_API_KEY", "abc")
+        # test non-existing env vars are not tracked
+        try:
+            os.environ["ANOTHER_API_KEY"]
+        except KeyError:
+            pass
+        assert all(x in tracked_env_names for x in ["DATABRICKS_HOST", "TEST_API_KEY"])
+        assert all(x not in tracked_env_names for x in ["INVALID_API_KEY", "ANOTHER_API_KEY"])
+
+    assert isinstance(os.environ, os._Environ)
+    assert all(x in os.environ for x in ["DATABRICKS_HOST", "TEST_API_KEY"])
+    assert all(x not in os.environ for x in ["INVALID_API_KEY", "ANOTHER_API_KEY"])
+
+    monkeypatch.setenv(MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING.name, "false")
+    with env_var_tracker() as env:
+        os.environ.get("API_KEY")
+        assert env == set()