1  import difflib
  2  import re
  3  from pathlib import Path
  4  
  5  import pytest
  6  from packaging.version import Version
  7  
  8  from dev.update_mlflow_versions import (
  9      get_current_py_version,
 10      replace_java,
 11      replace_java_pom_xml,
 12      replace_pyproject_toml,
 13      replace_python,
 14      replace_r,
 15      replace_ts,
 16  )
 17  
 18  # { filename: expected lines changed }
 19  _JAVA_FILES = {}
 20  
 21  _JAVA_XML_FILES = {
 22      "mlflow/java/pom.xml": {
 23          6: "  <version>{new_version}</version>",
 24          52: "    <mlflow-version>{new_version}</mlflow-version>",
 25      },
 26      "mlflow/java/client/pom.xml": {
 27          8: "    <version>{new_version}</version>",
 28      },
 29      "mlflow/java/spark_2.12/pom.xml": {
 30          4: "  <version>{new_version}</version>",
 31          18: "    <version>{new_version}</version>",
 32      },
 33      "mlflow/java/spark_2.13/pom.xml": {
 34          4: "  <version>{new_version}</version>",
 35          18: "    <version>{new_version}</version>",
 36      },
 37  }
 38  
 39  _TS_FILES = {
 40      "mlflow/server/js/src/common/constants.tsx": {
 41          12: "export const Version = '{new_version}';",
 42      },
 43      "docs/src/constants.ts": {
 44          1: "export const Version = '{new_version}';",
 45      },
 46  }
 47  
 48  _PYTHON_FILES = {
 49      "mlflow/version.py": {
 50          5: 'VERSION = "{new_version}"',
 51      }
 52  }
 53  
 54  _PYPROJECT_TOML_FILES = {
 55      "pyproject.toml": {
 56          12: 'version = "{new_version}"',
 57      },
 58      "pyproject.release.toml": {
 59          12: 'version = "{new_version}"',
 60          30: '  "mlflow-skinny=={new_version}",',
 61          31: '  "mlflow-tracing=={new_version}",',
 62      },
 63      "libs/skinny/pyproject.toml": {
 64          10: 'version = "{new_version}"',
 65      },
 66      "libs/tracing/pyproject.toml": {
 67          10: 'version = "{new_version}"',
 68      },
 69  }
 70  
 71  _R_FILES = {
 72      "mlflow/R/mlflow/DESCRIPTION": {
 73          4: "Version: {new_version}",
 74      }
 75  }
 76  
 77  _DIFF_REGEX = re.compile(r"--- (\d+,?\d*) ----")
 78  
 79  old_version = Version(get_current_py_version())
 80  _NEW_PY_VERSION = f"{old_version.major}.{old_version.minor}.{old_version.micro + 1}"
 81  
 82  
 83  def copy_and_run_change_func(monkeypatch, tmp_path, paths_to_copy, replace_func, new_version):
 84      for path in paths_to_copy:
 85          copy_path = tmp_path / path
 86          copy_path.parent.mkdir(parents=True, exist_ok=True)
 87          copy_path.write_text(path.read_text())
 88  
 89      with monkeypatch.context() as m:
 90          m.chdir(tmp_path)
 91  
 92          # pyproject.toml replace doesn't search for the old version,
 93          # it just replaces the version line with the new version.
 94          if replace_func == replace_pyproject_toml:
 95              replace_func(new_version, paths_to_copy)
 96          else:
 97              replace_func(str(old_version), new_version, paths_to_copy)
 98  
 99  
100  @pytest.mark.parametrize(
101      ("replace_func", "expect_dict", "new_py_version", "expected_new_version"),
102      [
103          (replace_java, _JAVA_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
104          (replace_java, _JAVA_FILES, _NEW_PY_VERSION + ".dev0", _NEW_PY_VERSION + "-SNAPSHOT"),
105          (replace_java, _JAVA_FILES, _NEW_PY_VERSION + "rc1", _NEW_PY_VERSION + "-SNAPSHOT"),
106          (replace_java_pom_xml, _JAVA_XML_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
107          (
108              replace_java_pom_xml,
109              _JAVA_XML_FILES,
110              _NEW_PY_VERSION + ".dev0",
111              _NEW_PY_VERSION + "-SNAPSHOT",
112          ),
113          (
114              replace_java_pom_xml,
115              _JAVA_XML_FILES,
116              _NEW_PY_VERSION + "rc1",
117              _NEW_PY_VERSION + "-SNAPSHOT",
118          ),
119          (replace_ts, _TS_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
120          (replace_python, _PYTHON_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
121          (replace_pyproject_toml, _PYPROJECT_TOML_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
122          (replace_r, _R_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
123      ],
124  )
125  def test_update_mlflow_versions(
126      monkeypatch, tmp_path, replace_func, expect_dict, new_py_version, expected_new_version
127  ):
128      paths_to_change = [Path(filename) for filename in expect_dict]
129      copy_and_run_change_func(
130          monkeypatch,
131          tmp_path,
132          # always copy version.py since we need it in get_current_py_version()
133          paths_to_change + [Path("mlflow/version.py")],
134          replace_func,
135          new_py_version,
136      )
137  
138      # diff files
139      for filename, expected_changes in expect_dict.items():
140          old_file = Path(filename).read_text().splitlines()
141          new_file = (tmp_path / filename).read_text().splitlines()
142          diff = list(difflib.context_diff(old_file, new_file, n=0))
143          changed_lines = _parse_diff_line(diff)
144  
145          formatted_expected_changes = {
146              line_num: change.format(new_version=expected_new_version)
147              for line_num, change in expected_changes.items()
148          }
149  
150          assert changed_lines == formatted_expected_changes
151  
152  
153  def _parse_diff_line(diff: list[str]) -> dict[int, str]:
154      diff_lines = {}
155      for idx, line in enumerate(diff):
156          match = _DIFF_REGEX.search(line)
157          if not match:
158              continue
159  
160          if "," in match.group(1):
161              # multi-line change is represented as [(start,end), line1, line2, ...]
162              start, end = map(int, match.group(1).split(","))
163              for i in range(start, end + 1):
164                  # the [2:] is to cut out the "! " at the beginning of diff lines
165                  diff_lines[i] = diff[idx + (i - start) + 1][2:]
166          else:
167              # single-line change
168              diff_lines[int(match.group(1))] = diff[idx + 1][2:]
169  
170      return diff_lines