From 63e416fc6b05cc6bfb8ab58abf2083878c3e3269 Mon Sep 17 00:00:00 2001 From: John Sirois Date: Fri, 1 Nov 2024 15:36:14 -0700 Subject: [PATCH] Start to fill out some tests... --- pex/cache/dirs.py | 8 + pex/cache/prunable.py | 9 +- pex/cli/commands/cache/command.py | 112 +++---- .../cli/commands/test_cache_prune.py | 285 ++++++++++++++++-- 4 files changed, 339 insertions(+), 75 deletions(-) diff --git a/pex/cache/dirs.py b/pex/cache/dirs.py index d5a4d99bf..bbfa3cdfc 100644 --- a/pex/cache/dirs.py +++ b/pex/cache/dirs.py @@ -229,6 +229,14 @@ def __init__( # type: (...) -> None self.path = path + def __eq__(self, other): + # type: (Any) -> bool + return type(self) == type(other) and self.path == other.path + + def __hash__(self): + # type: () -> int + return hash(self.path) + def __repr__(self): # type: () -> str return "{clazz}(path={path})".format(clazz=self.__class__.__name__, path=self.path) diff --git a/pex/cache/prunable.py b/pex/cache/prunable.py index d67969d5e..a517a230c 100644 --- a/pex/cache/prunable.py +++ b/pex/cache/prunable.py @@ -109,12 +109,15 @@ def scan(cls, cutoff): ) # type: Container[Union[BootstrapDir, UserCodeDir, InstalledWheelDir]] pips = attr.ib() # type: Pips - def iter_pex_deps(self): + def iter_pex_unused_deps(self): # type: () -> Iterator[Union[BootstrapDir, UserCodeDir, InstalledWheelDir]] - return iter(self._pex_deps) + for dep in self._pex_deps: + if dep not in self._unprunable_deps: + yield dep - def iter_unused_deps(self): + def iter_other_unused_deps(self): # type: () -> Iterator[Union[BootstrapDir, UserCodeDir, InstalledWheelDir]] + for bootstrap_dir in BootstrapDir.iter_all(): if bootstrap_dir not in self._pex_deps and bootstrap_dir not in self._unprunable_deps: yield bootstrap_dir diff --git a/pex/cli/commands/cache/command.py b/pex/cli/commands/cache/command.py index e7ac343a3..c6325e98b 100644 --- a/pex/cli/commands/cache/command.py +++ b/pex/cli/commands/cache/command.py @@ -98,6 +98,59 @@ def parse(cls, spec): cutoff = attr.ib() # type: datetime +def _prune_cache_dir( + dry_run, # type: bool + additional_cache_dirs_by_project_name_and_version, # type: Mapping[Tuple[ProjectName, Version], Iterable[AtomicCacheDir]] + cache_dir, # type: AtomicCacheDir +): + # type: (...) -> DiskUsage + paths_to_prune = [] # type: List[str] + + def prune_if_exists(path): + # type: (Optional[str]) -> None + if path and os.path.exists(path): + paths_to_prune.append(path) + + if isinstance(cache_dir, InstalledWheelDir): + paths_to_prune.append(os.path.dirname(cache_dir.path)) + prune_if_exists(CacheDir.PACKED_WHEELS.path(cache_dir.install_hash)) + for additional_dir in additional_cache_dirs_by_project_name_and_version.get( + (cache_dir.project_name, cache_dir.version), () + ): + prune_if_exists(additional_dir) + elif isinstance(cache_dir, BootstrapDir): + paths_to_prune.append(cache_dir.path) + prune_if_exists(CacheDir.BOOTSTRAP_ZIPS.path(cache_dir.bootstrap_hash)) + else: + paths_to_prune.append(cache_dir.path) + + disk_usages = [DiskUsage.collect(path) for path in paths_to_prune] + if not dry_run: + for path in paths_to_prune: + safe_rmtree(path) + if isinstance(cache_dir, InstalledWheelDir) and cache_dir.symlink_dir: + safe_rmtree(cache_dir.symlink_dir) + elif isinstance(cache_dir, VenvDirs): + safe_rmtree(cache_dir.short_dir) + + return ( + disk_usages[0] + if len(disk_usages) == 1 + else DiskUsage.aggregate(cache_dir.path, disk_usages) + ) + + +def _prune_pip( + dry_run, # type: bool + pip_path_to_prune, # type: str +): + # type: (...) -> DiskUsage + du = DiskUsage.collect(pip_path_to_prune) + if not dry_run: + safe_rmtree(pip_path_to_prune) + return du + + class Cache(OutputMixin, BuildTimeCommand): """Interact with the Pex cache.""" @@ -467,54 +520,6 @@ def _purge(self): return Ok() - def _prune_cache_dir( - self, - additional_cache_dirs_by_project_name_and_version, # type: Mapping[Tuple[ProjectName, Version], Iterable[AtomicCacheDir]] - cache_dir, # type: AtomicCacheDir - ): - # type: (...) -> DiskUsage - paths_to_prune = [] # type: List[str] - - def prune_if_exists(path): - # type: (Optional[str]) -> None - if path and os.path.exists(path): - paths_to_prune.append(path) - - if isinstance(cache_dir, InstalledWheelDir): - paths_to_prune.append(os.path.dirname(cache_dir.path)) - prune_if_exists(CacheDir.PACKED_WHEELS.path(cache_dir.install_hash)) - for additional_dir in additional_cache_dirs_by_project_name_and_version.get( - (cache_dir.project_name, cache_dir.version), () - ): - prune_if_exists(additional_dir) - elif isinstance(cache_dir, BootstrapDir): - paths_to_prune.append(cache_dir.path) - prune_if_exists(CacheDir.BOOTSTRAP_ZIPS.path(cache_dir.bootstrap_hash)) - else: - paths_to_prune.append(cache_dir.path) - - disk_usages = [DiskUsage.collect(path) for path in paths_to_prune] - if not self.options.dry_run: - for path in paths_to_prune: - safe_rmtree(path) - if isinstance(cache_dir, InstalledWheelDir) and cache_dir.symlink_dir: - safe_rmtree(cache_dir.symlink_dir) - elif isinstance(cache_dir, VenvDirs): - safe_rmtree(cache_dir.short_dir) - - return ( - disk_usages[0] - if len(disk_usages) == 1 - else DiskUsage.aggregate(cache_dir.path, disk_usages) - ) - - def _prune_pip(self, pip_path_to_prune): - # type: (str) -> DiskUsage - du = DiskUsage.collect(pip_path_to_prune) - if not self.options.dry_run: - safe_rmtree(pip_path_to_prune) - return du - def _prune(self): # type: () -> Result @@ -534,7 +539,7 @@ def _prune(self): cutoff = self.options.cutoff prunable = Prunable.scan(cutoff.cutoff) - unused_deps = tuple(prunable.iter_unused_deps()) + unused_deps = tuple(prunable.iter_other_unused_deps()) unused_wheels = tuple(dep for dep in unused_deps if isinstance(dep, InstalledWheelDir)) additional_cache_dirs_by_project_name_and_version = defaultdict( @@ -549,8 +554,11 @@ def _prune(self): ].append(cache_dir) prune_cache_dir = functools.partial( - self._prune_cache_dir, additional_cache_dirs_by_project_name_and_version + _prune_cache_dir, + self.options.dry_run, + additional_cache_dirs_by_project_name_and_version, ) + prune_pip = functools.partial(_prune_pip, self.options.dry_run) def prune_unused_deps(additional=False): # type: (bool) -> None @@ -596,7 +604,7 @@ def prune_pips(): tuple( iter_map_parallel( prunable.pips.paths, - function=self._prune_pip, + function=prune_pip, noun="Pip", verb="prune", verb_past="pruned", @@ -773,7 +781,7 @@ def prune_interpreters(): ) print(file=fp) - deps = tuple(prunable.iter_pex_deps()) + deps = tuple(prunable.iter_pex_unused_deps()) if self.options.dry_run: print( "Might have pruned up to {count} {cached_pex_dependency}.".format( diff --git a/tests/integration/cli/commands/test_cache_prune.py b/tests/integration/cli/commands/test_cache_prune.py index d9a372342..340fe1313 100644 --- a/tests/integration/cli/commands/test_cache_prune.py +++ b/tests/integration/cli/commands/test_cache_prune.py @@ -1,52 +1,297 @@ # Copyright 2024 Pex project contributors. # Licensed under the Apache License, Version 2.0 (see LICENSE). +from __future__ import absolute_import + +import os.path +import shutil +import subprocess +import time +from datetime import datetime, timedelta +from textwrap import dedent + +import colors +import pytest + +from pex.cache import access +from pex.cache.dirs import BootstrapDir, CacheDir, InstalledWheelDir, UnzipDir +from pex.cli.commands.cache.du import DiskUsage +from pex.common import safe_open +from pex.pep_503 import ProjectName +from pex.pex_info import PexInfo +from pex.pip.version import PipVersion from pex.typing import TYPE_CHECKING +from pex.variables import ENV +from testing import environment_as, make_env, run_pex_command +from testing.cli import run_pex3 +from testing.pytest.tmp import Tempdir if TYPE_CHECKING: - from typing import Any + from typing import Iterable, Iterator, Optional -def test_nothing_prunable(tmpdir): - # type: (Any) -> None - pass +@pytest.fixture +def pex_root(tmpdir): + # type: (Tempdir) -> Iterator[str] + _pex_root = tmpdir.join("pex_root") + with ENV.patch(PEX_ROOT=_pex_root) as env, environment_as(**env): + yield _pex_root -def test_installed_wheel_prune_build_time(tmpdir): - # type: (Any) -> None - pass +@pytest.fixture +def pex(tmpdir): + # type: (Tempdir) -> str + return tmpdir.join("pex") -def test_installed_wheel_prune_run_time(tmpdir): - # type: (Any) -> None - pass +@pytest.fixture +def lock(tmpdir): + # type: (Tempdir) -> str + return tmpdir.join("lock.json") -def test_zipapp_prune(tmpdir): - # type: (Any) -> None - pass +def test_nothing_prunable( + pex_root, # type: str + pex, # type: str +): + # type: (...) -> None + run_pex_command(args=["-o", pex]).assert_success() + pex_size = os.path.getsize(pex) + + subprocess.check_call(args=[pex, "-c", ""]) + pre_prune_du = DiskUsage.collect(pex_root) + assert ( + pre_prune_du.size > pex_size + ), "Expected the unzipped PEX to be larger than the zipped pex." + + # The default prune threshold should be high enough to never trigger in a test run (it's 2 + # weeks old at the time of writing). + run_pex3("cache", "prune").assert_success() + assert pre_prune_du == DiskUsage.collect(pex_root) + + +def test_installed_wheel_prune_build_time( + pex_root, # type: str + pex, # type: str +): + # type: (...) -> None + + run_pex_command(args=["ansicolors==1.1.8", "-o", pex]).assert_success() + installed_wheels_size = DiskUsage.collect(CacheDir.INSTALLED_WHEELS.path()).size + assert installed_wheels_size > 0 + assert 0 == DiskUsage.collect(CacheDir.UNZIPPED_PEXES.path()).size + assert 0 == DiskUsage.collect(CacheDir.BOOTSTRAPS.path()).size + assert 0 == DiskUsage.collect(CacheDir.USER_CODE.path()).size + + run_pex3("cache", "prune", "--older-than", "0 seconds").assert_success() + assert 0 == DiskUsage.collect(CacheDir.UNZIPPED_PEXES.path()).size + assert 0 == DiskUsage.collect(CacheDir.INSTALLED_WHEELS.path()).size + assert 0 == DiskUsage.collect(CacheDir.BOOTSTRAPS.path()).size + assert 0 == DiskUsage.collect(CacheDir.USER_CODE.path()).size + + +def test_installed_wheel_prune_run_time( + pex_root, # type: str + pex, # type: str +): + # type: (...) -> None + + run_pex_command(args=["cowsay==5.0", "-c", "cowsay", "-o", pex]).assert_success() + pex_size = os.path.getsize(pex) + + shutil.rmtree(pex_root) + assert 0 == DiskUsage.collect(pex_root).size + + assert b"| Moo! |" in subprocess.check_output(args=[pex, "Moo!"]) + pre_prune_du = DiskUsage.collect(pex_root) + assert DiskUsage.collect(CacheDir.INSTALLED_WHEELS.path()).size > 0 + assert DiskUsage.collect(CacheDir.UNZIPPED_PEXES.path()).size > 0 + assert DiskUsage.collect(CacheDir.BOOTSTRAPS.path()).size > 0 + assert DiskUsage.collect(CacheDir.USER_CODE.path()).size > 0 + assert ( + pre_prune_du.size > pex_size + ), "Expected the unzipped PEX to be larger than the zipped pex." + + run_pex3("cache", "prune", "--older-than", "0 seconds").assert_success() + assert 0 == DiskUsage.collect(CacheDir.UNZIPPED_PEXES.path()).size + assert 0 == DiskUsage.collect(CacheDir.INSTALLED_WHEELS.path()).size + assert 0 == DiskUsage.collect(CacheDir.BOOTSTRAPS.path()).size + assert 0 == DiskUsage.collect(CacheDir.USER_CODE.path()).size + + +def test_zipapp_prune( + pex_root, # type: str + pex, # type: str + tmpdir, # type: Tempdir +): + # type: (...) -> None + + with safe_open(tmpdir.join("src", "app.py"), "w") as fp: + fp.write( + dedent( + """\ + import colors + + + if __name__ == "__main__": + print(colors.green("Hello Cache!")) + """ + ) + ) + run_pex_command( + args=["ansicolors==1.1.8", "-D", "src", "-m" "app", "-o", pex], cwd=tmpdir.path + ).assert_success() + pex_size = os.path.getsize(pex) + installed_wheels_size = DiskUsage.collect(CacheDir.INSTALLED_WHEELS.path()).size + assert installed_wheels_size > 0 + assert 0 == DiskUsage.collect(CacheDir.UNZIPPED_PEXES.path()).size + assert 0 == DiskUsage.collect(CacheDir.BOOTSTRAPS.path()).size + assert 0 == DiskUsage.collect(CacheDir.USER_CODE.path()).size + + assert ( + colors.green("Hello Cache!") == subprocess.check_output(args=[pex]).decode("utf-8").strip() + ) + pre_prune_du = DiskUsage.collect(pex_root) + assert ( + DiskUsage.collect(CacheDir.INSTALLED_WHEELS.path()).size > installed_wheels_size + ), "Expected .pyc files to be compiled leading to more disk space usage" + assert DiskUsage.collect(CacheDir.UNZIPPED_PEXES.path()).size > 0 + assert DiskUsage.collect(CacheDir.BOOTSTRAPS.path()).size > 0 + assert DiskUsage.collect(CacheDir.USER_CODE.path()).size > 0 + assert ( + pre_prune_du.size > pex_size + ), "Expected the unzipped PEX to be larger than the zipped pex." + + run_pex3("cache", "prune", "--older-than", "0 seconds").assert_success() + assert 0 == DiskUsage.collect(CacheDir.UNZIPPED_PEXES.path()).size + assert 0 == DiskUsage.collect(CacheDir.INSTALLED_WHEELS.path()).size + assert 0 == DiskUsage.collect(CacheDir.BOOTSTRAPS.path()).size + assert 0 == DiskUsage.collect(CacheDir.USER_CODE.path()).size -def test_zipapp_prune_shared_bootstrap(tmpdir): - # type: (Any) -> None - pass + +def set_last_access_one_day_ago(pex): + # type: (str) -> None + + one_day_ago = time.mktime((datetime.now() - timedelta(days=1)).timetuple()) + pex_info = PexInfo.from_pex(pex) + if pex_info.venv: + pex_dir = pex_info.runtime_venv_dir(pex) + assert pex_dir is not None + access.record_access(pex_dir, one_day_ago) + else: + assert pex_info.pex_hash is not None + access.record_access(UnzipDir.create(pex_info.pex_hash), one_day_ago) + + +def assert_installed_wheels( + names, # type: Iterable[str] + message=None, # type: Optional[str] +): + expected = set(map(ProjectName, names)) + actual = {iwd.project_name for iwd in InstalledWheelDir.iter_all()} + if message: + assert expected == actual, message + else: + assert expected == actual + + +def expected_pip_wheels(): + # type: () -> Iterable[str] + if PipVersion.DEFAULT is PipVersion.VENDORED: + return "pip", "setuptools" + else: + return "pip", "setuptools", "wheel" + + +def expected_pip_wheels_plus(*names): + # type: (*str) -> Iterable[str] + wheels = list(expected_pip_wheels()) + wheels.extend(names) + return wheels + + +def test_zipapp_prune_shared_bootstrap( + pex_root, # type: str + pex, # type: str + tmpdir, # type: Tempdir +): + # type: (...) -> None + + with safe_open(tmpdir.join("src", "app.py"), "w") as fp: + fp.write( + dedent( + """\ + import colors + + + if __name__ == "__main__": + print(colors.green("Hello Cache!")) + """ + ) + ) + run_pex_command( + args=["ansicolors==1.1.8", "-D", "src", "-m" "app", "-o", pex], cwd=tmpdir.path + ).assert_success() + assert ( + colors.green("Hello Cache!") == subprocess.check_output(args=[pex]).decode("utf-8").strip() + ) + set_last_access_one_day_ago(pex) + + empty_pex = tmpdir.join("empty.pex") + run_pex_command(args=["-o", empty_pex]).assert_success() + subprocess.check_call(args=[empty_pex, "-c", ""]) + + bootstraps = list(BootstrapDir.iter_all()) + assert len(bootstraps) == 1, "Expected a shared bootstrap between pex and empty.pex." + bootstrap = bootstraps[0] + + assert_installed_wheels( + expected_pip_wheels_plus("ansicolors"), + ( + "There should be an ansicolors wheel for the pex as well as pip, setuptools and wheel wheels " + "for at least 1 Pip." + ), + ) + + run_pex3( + "cache", "prune", "--older-than", "1 hour", env=make_env(PEX_VERBOSE=1) + ).assert_success() + + assert [bootstrap] == list(BootstrapDir.iter_all()) + assert_installed_wheels(expected_pip_wheels()) def test_zipapp_prune_shared_code(tmpdir): - # type: (Any) -> None + # type: (Tempdir) -> None pass def test_zipapp_prune_shared_deps(tmpdir): - # type: (Any) -> None + # type: (Tempdir) -> None + pass + + +def test_venv_prune(tmpdir): + # type: (Tempdir) -> None + pass + + +def test_venv_prune_shared_deps(tmpdir): + # type: (Tempdir) -> None pass def test_venv_prune_symlinks(tmpdir): - # type: (Any) -> None + # type: (Tempdir) -> None pass def test_venv_prune_no_symlinks(tmpdir): - # type: (Any) -> None + # type: (Tempdir) -> None + pass + + +def test_venv_prune_interpreters(tmpdir): + # type: (Tempdir) -> None pass