1  from unittest import mock
  2  
  3  import pytest
  4  
  5  from mlflow.utils import PYTHON_VERSION
  6  from mlflow.utils.environment import _PythonEnv
  7  
  8  
  9  def test_constructor_argument_validation():
 10      with pytest.raises(TypeError, match="`python` must be a string"):
 11          _PythonEnv(python=1)
 12  
 13      with pytest.raises(TypeError, match="`build_dependencies` must be a list"):
 14          _PythonEnv(build_dependencies=0)
 15  
 16      with pytest.raises(TypeError, match="`dependencies` must be a list"):
 17          _PythonEnv(dependencies=0)
 18  
 19  
 20  def test_to_yaml(tmp_path):
 21      yaml_path = tmp_path / "python_env.yaml"
 22      _PythonEnv(PYTHON_VERSION, ["a"], ["b"]).to_yaml(yaml_path)
 23      expected_content = f"""
 24  python: {PYTHON_VERSION}
 25  build_dependencies:
 26  - a
 27  dependencies:
 28  - b
 29  """.lstrip()
 30      assert yaml_path.read_text() == expected_content
 31  
 32  
 33  def test_from_yaml(tmp_path):
 34      content = f"""
 35  python: {PYTHON_VERSION}
 36  build_dependencies:
 37  - a
 38  - b
 39  dependencies:
 40  - c
 41  - d
 42  """
 43      yaml_path = tmp_path / "test.yaml"
 44      yaml_path.write_text(content)
 45      python_env = _PythonEnv.from_yaml(yaml_path)
 46      assert python_env.python == PYTHON_VERSION
 47      assert python_env.build_dependencies == ["a", "b"]
 48      assert python_env.dependencies == ["c", "d"]
 49  
 50  
 51  def test_from_conda_yaml(tmp_path):
 52      content = f"""
 53  name: example
 54  channels:
 55    - conda-forge
 56  dependencies:
 57    - python={PYTHON_VERSION}
 58    - pip
 59    - pip:
 60      - a
 61      - b
 62  """
 63      yaml_path = tmp_path / "conda.yaml"
 64      yaml_path.write_text(content)
 65      python_env = _PythonEnv.from_conda_yaml(yaml_path)
 66      assert python_env.python == PYTHON_VERSION
 67      assert python_env.build_dependencies == ["pip"]
 68      assert python_env.dependencies == ["a", "b"]
 69  
 70  
 71  def test_from_conda_yaml_build_dependencies(tmp_path):
 72      content = f"""
 73  name: example
 74  channels:
 75    - conda-forge
 76  dependencies:
 77    - python={PYTHON_VERSION}
 78    - pip=1.2.3
 79    - wheel==4.5.6
 80    - setuptools<=7.8.9
 81    - pip:
 82      - a
 83      - b
 84  """
 85      yaml_path = tmp_path / "conda.yaml"
 86      yaml_path.write_text(content)
 87      python_env = _PythonEnv.from_conda_yaml(yaml_path)
 88      assert python_env.python == PYTHON_VERSION
 89      assert python_env.build_dependencies == ["pip==1.2.3", "wheel==4.5.6", "setuptools<=7.8.9"]
 90      assert python_env.dependencies == ["a", "b"]
 91  
 92  
 93  def test_from_conda_yaml_use_current_python_version_when_no_python_spec_in_conda_yaml(tmp_path):
 94      content = """
 95  name: example
 96  channels:
 97    - conda-forge
 98  dependencies:
 99    - pip
100    - pip:
101      - a
102      - b
103  """
104      yaml_path = tmp_path / "conda.yaml"
105      yaml_path.write_text(content)
106      assert _PythonEnv.from_conda_yaml(yaml_path).python == PYTHON_VERSION
107  
108  
109  def test_from_conda_yaml_invalid_python_comparator(tmp_path):
110      content = f"""
111  name: example
112  channels:
113    - conda-forge
114  dependencies:
115    - python<{PYTHON_VERSION}
116    - pip:
117      - a
118      - b
119  """
120      yaml_path = tmp_path / "conda.yaml"
121      yaml_path.write_text(content)
122      with pytest.raises(Exception, match="Invalid version comparator for python"):
123          _PythonEnv.from_conda_yaml(yaml_path)
124  
125  
126  def test_from_conda_yaml_conda_dependencies_warning(tmp_path):
127      content = f"""
128  name: example
129  channels:
130    - conda-forge
131  dependencies:
132    - python={PYTHON_VERSION}
133    - foo
134    - bar
135    - pip:
136      - a
137  """
138      yaml_path = tmp_path / "conda.yaml"
139      yaml_path.write_text(content)
140      with mock.patch("mlflow.utils.environment._logger.warning") as mock_warning:
141          _PythonEnv.from_conda_yaml(yaml_path)
142          mock_warning.assert_called_with(
143              "The following conda dependencies will not be installed "
144              "in the resulting environment: %s",
145              ["foo", "bar"],
146          )