mlflow--mlflow/dev/flavors/tests/test_update.py

from datetime import datetime, timedelta, timezone
from pathlib import Path
from unittest import mock

import pytest
from flavors import _update
from flavors._update import VersionInfo
from pypi import Package


def _iso8601(dt: datetime) -> str:
    return (dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)).isoformat()


def package_from_version_infos(version_infos: list[VersionInfo]) -> Package:
    return Package.from_json({
        "releases": {
            v.version: [
                {
                    "filename": v.version + ".whl",
                    "upload_time_iso_8601": _iso8601(v.upload_time),
                }
            ]
            for v in version_infos
        },
    })


def package_from_versions(versions: list[str]) -> Package:
    fixed = datetime(2023, 10, 4, 16, 38, 57, tzinfo=timezone.utc)
    return package_from_version_infos([VersionInfo(v, fixed) for v in versions])


@pytest.fixture(autouse=True)
def change_working_directory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    """
    Changes the current working directory to a temporary directory to avoid modifying files in the
    repository.
    """
    monkeypatch.chdir(tmp_path)


def run_test(src: str, src_expected: str, mock_packages: dict[str, Package]) -> None:
    async def fake_get_packages(names: list[str]) -> list[Package]:
        return [mock_packages[n] for n in names]

    versions_yaml = Path("mlflow/ml-package-versions.yml")
    versions_yaml.parent.mkdir()
    versions_yaml.write_text(src)

    with (
        mock.patch("flavors._update.get_packages", new=fake_get_packages),
        mock.patch("flavors._update.check_pypi_accessibility") as mock_check,
    ):
        _update.update()
        mock_check.assert_called_once()

    assert versions_yaml.read_text() == src_expected


def test_multiple_flavors_are_correctly_updated() -> None:
    src = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    maximum: "0.0.1"
xgboost:
  package_info:
    pip_release: xgboost
  autologging:
    maximum: "0.1.1"
"""
    mock_packages = {
        "sklearn": package_from_versions(["0.0.2"]),
        "xgboost": package_from_versions(["0.1.2"]),
    }
    src_expected = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    maximum: "0.0.2"
xgboost:
  package_info:
    pip_release: xgboost
  autologging:
    maximum: "0.1.2"
"""
    run_test(src, src_expected, mock_packages)


def test_both_models_and_autologging_are_updated() -> None:
    src = """
sklearn:
  package_info:
    pip_release: sklearn
  models:
    maximum: "0.0.1"
  autologging:
    maximum: "0.0.1"
"""
    mock_packages = {
        "sklearn": package_from_versions(["0.0.2"]),
    }
    src_expected = """
sklearn:
  package_info:
    pip_release: sklearn
  models:
    maximum: "0.0.2"
  autologging:
    maximum: "0.0.2"
"""
    run_test(src, src_expected, mock_packages)


def test_pre_and_dev_versions_are_ignored() -> None:
    src = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    maximum: "0.0.1"
"""
    mock_packages = {
        "sklearn": package_from_versions([
            # pre-release and dev-release should be filtered out
            "0.0.3.rc1",  # pre-release
            "0.0.3.dev1",  # dev-release
            "0.0.2.post",  # post-release
            "0.0.2",  # final release
        ]),
    }
    src_expected = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    maximum: "0.0.2.post0"
"""
    run_test(src, src_expected, mock_packages)


def test_unsupported_versions_are_ignored() -> None:
    src = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    unsupported: ["0.0.3"]
    maximum: "0.0.1"
"""
    mock_packages = {"sklearn": package_from_versions(["0.0.2", "0.0.3"])}
    src_expected = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    unsupported: ["0.0.3"]
    maximum: "0.0.2"
"""
    run_test(src, src_expected, mock_packages)


def test_freeze_field_prevents_updating_maximum_version() -> None:
    src = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    pin_maximum: True
    maximum: "0.0.1"
"""
    mock_packages = {"sklearn": package_from_versions(["0.0.2"])}
    src_expected = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    pin_maximum: True
    maximum: "0.0.1"
"""
    run_test(src, src_expected, mock_packages)


def test_update_min_supported_version() -> None:
    src = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    minimum: "0.0.1"
    maximum: "0.0.8"
"""
    mock_packages = {
        "sklearn": package_from_version_infos([
            VersionInfo("0.0.2", datetime.now() - timedelta(days=1000)),
            VersionInfo("0.0.3", datetime.now() - timedelta(days=365)),
            VersionInfo("0.0.8", datetime.now() - timedelta(days=180)),
        ])
    }
    src_expected = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    minimum: "0.0.3"
    maximum: "0.0.8"
"""
    run_test(src, src_expected, mock_packages)


def test_update_min_supported_version_for_dead_package() -> None:
    src = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    minimum: "0.0.7"
    maximum: "0.0.8"
"""
    mock_packages = {
        "sklearn": package_from_version_infos([
            VersionInfo("0.0.7", datetime.now() - timedelta(days=1000)),
            VersionInfo("0.0.8", datetime.now() - timedelta(days=800)),
        ])
    }
    src_expected = """
sklearn:
  package_info:
    pip_release: sklearn
  autologging:
    minimum: "0.0.8"
    maximum: "0.0.8"
"""
    run_test(src, src_expected, mock_packages)