Improve version matching

This is still far from perfect, but it attempts a much more thorough comparison of the version reported by cog's CLI binary and the version reported by the embedded Python wheel.
replicate · Oct 18, 2024 · 72b2933 · 72b2933
1 parent 54eb576
commit 72b2933
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 14 deletions.
diff --git a/test-integration/test_integration/test_build.py b/test-integration/test_integration/test_build.py
@@ -1,11 +1,12 @@
 import json
 import os
 import subprocess
-import sys
 from pathlib import Path
 
 import pytest
 
+from .util import assert_versions_match
+
 
 def test_build_without_predictor(docker_image):
     project_dir = Path(__file__).parent / "fixtures/no-predictor-project"
@@ -318,12 +319,8 @@ def test_precompile(docker_image):
     assert build_process.returncode == 0
 
 
-@pytest.mark.skipif(
-    not sys.platform.startswith("linux"),
-    reason="only runs on linux due to CUDA binaries",
-)
 def test_cog_install_base_image(docker_image):
-    project_dir = Path(__file__).parent / "fixtures/torch-cuda-baseimage-project"
+    project_dir = Path(__file__).parent / "fixtures/string-project"
     build_process = subprocess.run(
         [
             "cog",
@@ -350,10 +347,7 @@ def test_cog_install_base_image(docker_image):
         capture_output=True,
     )
     assert cog_installed_version_process.returncode == 0
-    # Clean up the cog python version to go from 0.11.2.dev15+g54c08f0 to 0.11.2
-    cog_installed_version_stdout = ".".join(
-        cog_installed_version_process.stdout.decode().strip().split(".")[:3]
-    )
+    cog_installed_version = cog_installed_version_process.stdout.decode().strip()
     cog_version_process = subprocess.run(
         [
             "cog",
@@ -362,8 +356,9 @@ def test_cog_install_base_image(docker_image):
         cwd=project_dir,
         capture_output=True,
     )
-    cog_version_stdout = cog_version_process.stdout.decode().strip().split()[2]
-    # Clean up the cog go version to go from 0.11.2-dev+g54c08f0 to 0.11.2
-    cog_version_stdout = cog_version_stdout.split("-")[0]
+    cog_version = cog_version_process.stdout.decode().strip().split()[2]
 
-    assert cog_version_stdout == cog_installed_version_stdout
+    assert_versions_match(
+        semver_version=cog_version,
+        pep440_version=cog_installed_version,
+    )
diff --git a/test-integration/test_integration/util.py b/test-integration/test_integration/util.py
@@ -1,8 +1,97 @@
 import random
+import re
 import string
 import subprocess
 import time
 
+from packaging.version import VERSION_PATTERN
+
+# From the SemVer spec: https://semver.org/
+SEMVER_PATTERN = r"^(?P<major>0|[1-9]\d*)\.(?P<minor>0|[1-9]\d*)\.(?P<patch>0|[1-9]\d*)(?:-(?P<prerelease>(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+(?P<buildmetadata>[0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$"
+
+
+# Used to help ensure that the cog binary reports a semver version that matches
+# the PEP440 version of the embedded Python package.
+#
+# These are all valid pairs:
+#
+#   SEMVER            PEP440           NOTES
+#   0.11.2            0.11.2
+#   0.11.2-alpha2     0.11.2a2         prerelease counters are not checked
+#   0.11.2-beta1      0.11.2b1             "          "      "   "     "
+#   0.11.2-rc4        0.11.2rc4            "          "      "   "     "
+#   0.11.2-dev        0.11.2rc4.dev10  dev status overrides prerelease status
+#   0.11.2+gabcd      0.11.2+gabce
+#
+# The following are not valid pairs:
+#
+#   SEMVER            PEP440           NOTES
+#   0.11.2            0.11.3           mismatched release versions
+#   0.11.2-alpha2     0.11.2alpha2     PEP440 uses 'a' instead of 'alpha'
+#   0.11.2-alpha2     0.11.2b2         mismatched prerelease status
+#   0.11.2-rc4        0.11.2rc4.dev10  dev status should have overridden prerelease status
+#   0.11.2+gabcd      0.11.2+gdefg     mismatched local/build metadata
+#
+def assert_versions_match(semver_version: str, pep440_version: str):
+    semver_re = re.compile(SEMVER_PATTERN)
+    pep440_re = re.compile(VERSION_PATTERN, re.VERBOSE | re.IGNORECASE)
+
+    semver_match = semver_re.match(semver_version)
+    pep440_match = pep440_re.match(pep440_version)
+
+    assert semver_match, f"Invalid semver version: {semver_version}"
+    assert pep440_match, f"Invalid PEP 440 version: {pep440_version}"
+
+    semver_groups = semver_match.groupdict()
+    pep440_groups = pep440_match.groupdict()
+
+    semver_release = (
+        f"{semver_groups['major']}.{semver_groups['minor']}.{semver_groups['patch']}"
+    )
+
+    # Check base release version
+    assert (
+        semver_release == pep440_groups["release"]
+    ), f"Release versions do not match: {semver_release} != {pep440_groups['release']}"
+
+    # Check prerelease status
+    semver_pre = semver_groups["prerelease"]
+    pep440_pre = pep440_groups["pre"] or pep440_groups["dev"]
+
+    assert bool(semver_pre) == bool(pep440_pre), "Pre-release status does not match"
+
+    if semver_pre:
+        if semver_pre.startswith("alpha"):
+            assert (
+                pep440_groups["pre_l"] == "a"
+            ), "Alpha pre-release status does not match"
+            assert not pep440_groups[
+                "dev"
+            ], "Semver pre-release cannot also be a PEP440 dev build"
+
+        if semver_pre.startswith("beta"):
+            assert (
+                pep440_groups["pre_l"] == "b"
+            ), "Beta pre-release status does not match"
+            assert not pep440_groups[
+                "dev"
+            ], "Semver pre-release cannot also be a PEP440 dev build"
+
+        if semver_pre.startswith("rc"):
+            assert (
+                pep440_groups["pre_l"] == "rc"
+            ), "Release candidate pre-release status does not match"
+            assert not pep440_groups[
+                "dev"
+            ], "Semver pre-release cannot also be a PEP440 dev build"
+
+        if semver_pre.startswith("dev"):
+            assert pep440_groups["dev_l"] == "dev", "Dev build status does not match"
+
+    assert (
+        semver_groups["buildmetadata"] == pep440_groups["local"]
+    ), f"Local/build metadata component does not match: {semver_groups['buildmetadata']} != {pep440_groups['local']}"
+
 
 def random_string(length):
     return "".join(random.choice(string.ascii_lowercase) for i in range(length))

diff --git a/tox.ini b/tox.ini
@@ -65,6 +65,7 @@ base_python = python3.12
 changedir = test-integration
 skip_install = true
 deps =
+  packaging
   pytest
   pytest-rerunfailures
   pytest-timeout