Allow cluster_pack to handle several unzipped pexes in the same directory

cluster_pack/packaging.py

@@ -406,10 +406,16 @@ def detect_archive_names(
 
 
 def resolve_zip_from_pex_dir(pex_dir: str) -> str:
-    pex_files = glob.glob(f"{os.path.dirname(pex_dir)}/*.pex.zip")
-    assert len(pex_files) == 1, \
-        f"Expected to find single zipped PEX in same dir as {pex_dir}, got {pex_files}"
-    return pex_files[0]
+    parent_dir = os.path.dirname(pex_dir)
+    pex_files = glob.glob(f"{parent_dir}/*.pex.zip")
+    if len(pex_files) == 1:
+        return pex_files[0]
+
+    pex_file = next((x for x in pex_files if pex_dir == x.split(".zip")[0]), None)
+    if pex_file is None:
+        raise ValueError(f"{pex_dir}.zip not found, found {pex_files}")
+
+    return pex_file
 
 
 def detect_packer_from_spec(spec_file: str) -> Packer:

tests/test_packaging.py

@@ -12,7 +12,8 @@
 import pytest
 
 from cluster_pack import packaging, get_pyenv_usage_from_archive, uploader
-from cluster_pack.packaging import CONDA_CMD, UNPACKED_ENV_NAME, LARGE_PEX_CMD
+from cluster_pack.packaging import CONDA_CMD, UNPACKED_ENV_NAME, LARGE_PEX_CMD, \
+    resolve_zip_from_pex_dir
 
 MODULE_TO_TEST = "cluster_pack.packaging"
 MYARCHIVE_FILENAME = "myarchive.pex"
@@ -454,3 +455,64 @@ def test_detect_archive_names(running_from_pex, package_path, allow_large_pex, i
         actual, _, _ = packaging.detect_archive_names(
             packaging.PEX_PACKER, package_path, allow_large_pex)
         assert actual == expected
+
+
+def test_resolve_zip_from_pex_dir_with_1_pex():
+    # the function should work even if file and dir names are different if we handle only 1 pex
+    with tempfile.TemporaryDirectory() as tempdir:
+        open(os.path.join(tempdir, "pex1.pex.zip"), 'a').close()
+        pex_dir = os.path.join(tempdir, "pex_dir")
+        os.mkdir(pex_dir)
+
+        resolved = resolve_zip_from_pex_dir(pex_dir)
+
+    assert resolved == os.path.join(tempdir, 'pex1.pex.zip')
+
+
+def test_resolve_zip_from_pex_dir_with_2_pexes_with_correct_names():
+    with tempfile.TemporaryDirectory() as tempdir:
+        open(os.path.join(tempdir, "pex1.pex.zip"), 'a').close()
+        pex1_dir = os.path.join(tempdir, "pex1.pex")
+        os.mkdir(pex1_dir)
+
+        open(os.path.join(tempdir, "pex2.pex.zip"), 'a').close()
+        pex2_dir = os.path.join(tempdir, "pex2.pex")
+        os.mkdir(pex2_dir)
+
+        assert resolve_zip_from_pex_dir(pex1_dir) == os.path.join(tempdir, 'pex1.pex.zip')
+        assert resolve_zip_from_pex_dir(pex2_dir) == os.path.join(tempdir, 'pex2.pex.zip')
+
+
+def test_resolve_zip_from_pex_dir_with_2_pexes_with_overlapping_names():
+    with tempfile.TemporaryDirectory() as tempdir:
+        open(os.path.join(tempdir, "pex1.pex.zip"), 'a').close()
+        pex1_dir = os.path.join(tempdir, "pex1.pex")
+        os.mkdir(pex1_dir)
+
+        open(os.path.join(tempdir, "other_pex1.pex.zip"), 'a').close()
+        pex2_dir = os.path.join(tempdir, "other_pex1.pex")
+        os.mkdir(pex2_dir)
+
+        assert resolve_zip_from_pex_dir(pex1_dir) == os.path.join(tempdir, 'pex1.pex.zip')
+        assert resolve_zip_from_pex_dir(pex2_dir) == os.path.join(tempdir, 'other_pex1.pex.zip')
+
+
+def test_resolve_zip_from_pex_dir_with_2_pexes_with_wrong_names():
+    with tempfile.TemporaryDirectory() as tempdir:
+        open(os.path.join(tempdir, "pex1.pex.zip"), 'a').close()
+        pex1_dir = os.path.join(tempdir, "pex_dir.pex")
+        os.mkdir(pex1_dir)
+
+        open(os.path.join(tempdir, "pex2.pex.zip"), 'a').close()
+
+        with pytest.raises(ValueError, match=r".zip not found"):
+            resolve_zip_from_pex_dir(pex1_dir)
+
+
+def test_resolve_zip_from_pex_dir_with_no_zip_found():
+    with tempfile.TemporaryDirectory() as tempdir:
+        pex1_dir = os.path.join(tempdir, "pex_dir.pex")
+        os.mkdir(pex1_dir)
+
+        with pytest.raises(ValueError, match=r".zip not found"):
+            resolve_zip_from_pex_dir(pex1_dir)