test: add unit tests for nested_zip_loader

Author: ajaxbits

tests/test_nested_zip_loader.py

@@ -0,0 +1,88 @@
+import importlib
+import multiprocessing
+import shutil
+import sys
+import tempfile
+import zipfile
+from pathlib import Path
+
+import pytest
+
+# We have to use importlib here because nested_zip_loader calls load_nested_zip
+# at IMPORT TIME, which causes us a world of hurt in these tests if we try to
+# import it "normally" here.
+LOADER_PATH = Path(__file__).parent.parent / "package_python_function" / "nested_zip_loader.py"
+PKG_NAME = "_test_nested_zip"
+
+def _make_deps_zip(path: Path) -> None:
+    with zipfile.ZipFile(path, "w") as zf:
+        zf.writestr(f"{PKG_NAME}/__init__.py", "LOADED = True\n")
+
+@pytest.fixture()
+def lambda_env(tmp_path, monkeypatch):
+    """Simulate a Lambda-like layout: a task dir with a package whose __init__.py
+    is the nested_zip_loader code, and a .dependencies.zip with the 'real' code."""
+    task_dir = tmp_path / "task"
+    pkg_dir = task_dir / PKG_NAME
+    pkg_dir.mkdir(parents=True)
+    shutil.copy(LOADER_PATH, pkg_dir / "__init__.py")
+    _make_deps_zip(pkg_dir / ".dependencies.zip")
+
+    tmp_dir = tmp_path / "tmp"
+    tmp_dir.mkdir()
+    monkeypatch.setenv("TMPDIR", str(tmp_dir))
+    tempfile.tempdir = None
+
+    monkeypatch.syspath_prepend(str(task_dir))
+
+    yield tmp_path
+
+    sys.modules.pop(PKG_NAME, None)
+    tempfile.tempdir = None
+
+def test_cold_start_extracts(lambda_env):
+    mod = importlib.import_module(PKG_NAME)
+    assert mod.LOADED is True
+    assert (lambda_env / "tmp" / "package-python-function").exists()
+
+def test_warm_start_skips_extraction(lambda_env):
+    target_pkg = lambda_env / "tmp" / "package-python-function" / PKG_NAME
+    target_pkg.mkdir(parents=True)
+    (target_pkg / "__init__.py").write_text("LOADED = 'warm'\n")
+
+    mod = importlib.import_module(PKG_NAME)
+    assert mod.LOADED == "warm"
+
+def test_stale_staging_cleaned(lambda_env):
+    staging = lambda_env / "tmp" / ".stage.package-python-function"
+    staging.mkdir(parents=True)
+    (staging / "stale.txt").write_text("leftover")
+
+    importlib.import_module(PKG_NAME)
+    assert not staging.exists()
+
+def _fork_worker(results, index):
+    import importlib
+    import sys
+
+    sys.modules.pop(PKG_NAME, None)
+    try:
+        mod = importlib.import_module(PKG_NAME)
+        results[index] = mod.LOADED is True
+    except Exception as e:
+        print(f"Worker {index} failed: {e}")
+        results[index] = False
+
+def test_concurrent_no_race(lambda_env):
+    """Two forked processes race to extract; both should succeed."""
+    ctx = multiprocessing.get_context("fork")
+    results = ctx.Array("b", [0, 0])
+
+    procs = [ctx.Process(target=_fork_worker, args=(results, i)) for i in range(2)]
+    for p in procs:
+        p.start()
+    for p in procs:
+        p.join(timeout=10)
+
+    assert all(results), f"Workers returned: {list(results)}"
+    assert (lambda_env / "tmp" / "package-python-function").exists()