Source code for pex.testing

# Copyright 2014 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import absolute_import, print_function

import contextlib
import os
import platform
import random
import subprocess
import sys
from contextlib import contextmanager
from textwrap import dedent

from pex.common import open_zip, safe_mkdir, safe_mkdtemp, safe_rmtree, temporary_dir, touch
from pex.compatibility import to_unicode
from pex.distribution_target import DistributionTarget
from pex.executor import Executor
from pex.interpreter import PythonInterpreter
from pex.pex import PEX
from pex.pex_builder import PEXBuilder
from pex.pex_info import PexInfo
from pex.pip import get_pip
from pex.third_party.pkg_resources import Distribution
from pex.typing import TYPE_CHECKING
from pex.util import DistributionHelper, named_temporary_file

if TYPE_CHECKING:
    from typing import (
        Any,
        Callable,
        Dict,
        Iterable,
        Iterator,
        List,
        Mapping,
        Optional,
        Set,
        Text,
        Tuple,
        Union,
    )

PY_VER = sys.version_info[:2]
IS_PYPY = hasattr(sys, "pypy_version_info")
NOT_CPYTHON27 = IS_PYPY or PY_VER != (2, 7)
NOT_CPYTHON36 = IS_PYPY or PY_VER != (3, 6)
IS_LINUX = platform.system() == "Linux"
IS_NOT_LINUX = not IS_LINUX
NOT_CPYTHON27_OR_OSX = NOT_CPYTHON27 or IS_NOT_LINUX
NOT_CPYTHON36_OR_LINUX = NOT_CPYTHON36 or IS_LINUX


[docs]@contextlib.contextmanager def temporary_filename(): # type: () -> Iterator[str] """Creates a temporary filename. This is useful when you need to pass a filename to an API. Windows requires all handles to a file be closed before deleting/renaming it, so this makes it a bit simpler. """ with named_temporary_file() as fp: fp.write(b"") fp.close() yield fp.name
def random_bytes(length): # type: (int) -> bytes return "".join(map(chr, (random.randint(ord("a"), ord("z")) for _ in range(length)))).encode( "utf-8" )
[docs]def get_dep_dist_names_from_pex(pex_path, match_prefix=""): # type: (str, str) -> Set[str] """Given an on-disk pex, extract all of the unique first-level paths under `.deps`.""" with open_zip(pex_path) as pex_zip: dep_gen = (f.split(os.sep)[1] for f in pex_zip.namelist() if f.startswith(".deps/")) return set(item for item in dep_gen if item.startswith(match_prefix))
[docs]@contextlib.contextmanager def temporary_content(content_map, interp=None, seed=31337, perms=0o644): # type: (Mapping[str, Union[int, str]], Optional[Dict[str, Any]], int, int) -> Iterator[str] """Write content to disk where content is map from string => (int, string). If target is int, write int random bytes. Otherwise write contents of string. """ random.seed(seed) interp = interp or {} with temporary_dir() as td: for filename, size_or_content in content_map.items(): dest = os.path.join(td, filename) safe_mkdir(os.path.dirname(dest)) with open(dest, "wb") as fp: if isinstance(size_or_content, int): fp.write(random_bytes(size_or_content)) else: fp.write((size_or_content % interp).encode("utf-8")) os.chmod(dest, perms) yield td
@contextlib.contextmanager def make_project( name="my_project", # type: str version="0.0.0", # type: str zip_safe=True, # type: bool install_reqs=None, # type: Optional[List[str]] extras_require=None, # type: Optional[Dict[str, List[str]]] entry_points=None, # type: Optional[Union[str, Dict[str, List[str]]]] ): # type: (...) -> Iterator[str] project_content = { "setup.py": dedent( """ from setuptools import setup setup( name=%(project_name)r, version=%(version)r, zip_safe=%(zip_safe)r, packages=[%(project_name)r], scripts=[ 'scripts/hello_world', 'scripts/shell_script', ], package_data={%(project_name)r: ['package_data/*.dat']}, install_requires=%(install_requires)r, extras_require=%(extras_require)r, entry_points=%(entry_points)r, ) """ ), "scripts/hello_world": '#!/usr/bin/env python\nprint("hello world!")\n', "scripts/shell_script": "#!/usr/bin/env bash\necho hello world\n", os.path.join(name, "__init__.py"): 0, os.path.join(name, "my_module.py"): 'def do_something():\n print("hello world!")\n', os.path.join(name, "package_data/resource1.dat"): 1000, os.path.join(name, "package_data/resource2.dat"): 1000, } # type: Dict[str, Union[str, int]] interp = { "project_name": name, "version": version, "zip_safe": zip_safe, "install_requires": install_reqs or [], "extras_require": extras_require or {}, "entry_points": entry_points or {}, } with temporary_content(project_content, interp=interp) as td: yield td
[docs]class WheelBuilder(object): """Create a wheel distribution from an unpacked setup.py-based project."""
[docs] class BuildFailure(Exception): pass
def __init__(self, source_dir, interpreter=None, wheel_dir=None): # type: (str, Optional[PythonInterpreter], Optional[str]) -> None """Create a wheel from an unpacked source distribution in source_dir.""" self._source_dir = source_dir self._wheel_dir = wheel_dir or safe_mkdtemp() self._interpreter = interpreter or PythonInterpreter.get() def bdist(self): # type: () -> str get_pip().spawn_build_wheels( distributions=[self._source_dir], wheel_dir=self._wheel_dir, interpreter=self._interpreter, ).wait() dists = os.listdir(self._wheel_dir) if len(dists) == 0: raise self.BuildFailure("No distributions were produced!") if len(dists) > 1: raise self.BuildFailure("Ambiguous source distributions found: %s" % (" ".join(dists))) return os.path.join(self._wheel_dir, dists[0])
@contextlib.contextmanager def built_wheel( name="my_project", # type: str version="0.0.0", # type: str zip_safe=True, # type: bool install_reqs=None, # type: Optional[List[str]] extras_require=None, # type: Optional[Dict[str, List[str]]] interpreter=None, # type: Optional[PythonInterpreter] **kwargs # type: Any ): # type: (...) -> Iterator[str] with make_project( name=name, version=version, zip_safe=zip_safe, install_reqs=install_reqs, extras_require=extras_require, ) as td: builder = WheelBuilder(td, interpreter=interpreter, **kwargs) yield builder.bdist() @contextlib.contextmanager def make_source_dir( name="my_project", # type: str version="0.0.0", # type: str install_reqs=None, # type: Optional[List[str]] extras_require=None, # type: Optional[Dict[str, List[str]]] ): # type: (...) -> Iterator[str] with make_project( name=name, version=version, install_reqs=install_reqs, extras_require=extras_require ) as td: yield td @contextlib.contextmanager def make_bdist( name="my_project", # type: str version="0.0.0", # type: str zip_safe=True, # type: bool interpreter=None, # type: Optional[PythonInterpreter] **kwargs # type: Any ): # type: (...) -> Iterator[Distribution] with built_wheel( name=name, version=version, zip_safe=zip_safe, interpreter=interpreter, **kwargs ) as dist_location: install_dir = os.path.join(safe_mkdtemp(), os.path.basename(dist_location)) get_pip().spawn_install_wheel( wheel=dist_location, install_dir=install_dir, target=DistributionTarget.for_interpreter(interpreter), ).wait() dist = DistributionHelper.distribution_from_path(install_dir) assert dist is not None yield dist COVERAGE_PREAMBLE = """ try: from coverage import coverage cov = coverage(auto_data=True, data_suffix=True) cov.start() except ImportError: pass """
[docs]def write_simple_pex( td, # type: str exe_contents=None, # type: Optional[str] dists=None, # type: Optional[Iterable[Distribution]] sources=None, # type: Optional[Iterable[Tuple[str, str]]] coverage=False, # type: bool interpreter=None, # type: Optional[PythonInterpreter] pex_info=None, # type: Optional[PexInfo] ): # type: (...) -> PEXBuilder """Write a pex file that optionally contains an executable entry point. :param td: temporary directory path :param exe_contents: entry point python file :param dists: distributions to include, typically sdists or bdists :param sources: sources to include, as a list of pairs (env_filename, contents) :param coverage: include coverage header :param interpreter: a custom interpreter to use to build the pex :param pex_info: a custom PexInfo to use to build the pex. """ dists = dists or [] sources = sources or [] safe_mkdir(td) pb = PEXBuilder( path=td, preamble=COVERAGE_PREAMBLE if coverage else None, interpreter=interpreter, pex_info=pex_info, ) for dist in dists: pb.add_dist_location(dist.location if isinstance(dist, Distribution) else dist) for env_filename, contents in sources: src_path = os.path.join(td, env_filename) safe_mkdir(os.path.dirname(src_path)) with open(src_path, "w") as fp: fp.write(contents) pb.add_source(src_path, env_filename) if exe_contents: with open(os.path.join(td, "exe.py"), "w") as fp: fp.write(exe_contents) pb.set_executable(os.path.join(td, "exe.py")) pb.freeze() return pb
# TODO(#1041): use `typing.NamedTuple` once we require Python 3.
[docs]class IntegResults(object): """Convenience object to return integration run results.""" def __init__(self, output, error, return_code): # type: (Text, Text, int) -> None super(IntegResults, self).__init__() self.output = output self.error = error self.return_code = return_code def assert_success(self): # type: () -> None assert ( self.return_code == 0 ), "integration test failed: return_code={}, output={}, error={}".format( self.return_code, self.output, self.error ) def assert_failure(self): # type: () -> None assert self.return_code != 0
[docs]def run_pex_command(args, env=None, python=None, quiet=False): # type: (Iterable[str], Optional[Dict[str, str]], Optional[str], bool) -> IntegResults """Simulate running pex command for integration testing. This is different from run_simple_pex in that it calls the pex command rather than running a generated pex. This is useful for testing end to end runs with specific command line arguments or env options. """ cmd = [python or sys.executable, "-mpex"] if not quiet: cmd.append("-vvvvv") cmd.extend(args) process = Executor.open_process( cmd=cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE ) output, error = process.communicate() return IntegResults(output.decode("utf-8"), error.decode("utf-8"), process.returncode)
def run_simple_pex( pex, # type: str args=(), # type: Iterable[str] interpreter=None, # type: Optional[PythonInterpreter] stdin=None, # type: Optional[bytes] **kwargs # type: Any ): # type: (...) -> Tuple[bytes, int] p = PEX(pex, interpreter=interpreter) process = p.run( args=args, blocking=False, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, **kwargs ) stdout, _ = process.communicate(input=stdin) return stdout.replace(b"\r", b""), process.returncode def run_simple_pex_test( body, # type: str args=(), # type: Iterable[str] env=None, # type: Optional[Mapping[str, str]] dists=None, # type: Optional[Iterable[Distribution]] coverage=False, # type: bool interpreter=None, # type: Optional[PythonInterpreter] ): # type: (...) -> Tuple[bytes, int] with temporary_dir() as td1, temporary_dir() as td2: pb = write_simple_pex(td1, body, dists=dists, coverage=coverage, interpreter=interpreter) pex = os.path.join(td2, "app.pex") pb.build(pex) return run_simple_pex(pex, args=args, env=env, interpreter=interpreter) def bootstrap_python_installer(dest): # type: (str) -> None safe_rmtree(dest) for _ in range(3): try: subprocess.check_call(["git", "clone", "https://github.com/pyenv/pyenv.git", dest]) except subprocess.CalledProcessError as e: print("caught exception: %r" % e) continue else: break else: raise RuntimeError("Helper method could not clone pyenv from git after 3 tries") # Create an empty file indicating the fingerprint of the correct set of test interpreters. touch(os.path.join(dest, _INTERPRETER_SET_FINGERPRINT)) # NB: We keep the pool of bootstrapped interpreters as small as possible to avoid timeouts in CI # otherwise encountered when fetching and building too many on a cache miss. In the past we had # issues with the combination of 7 total unique interpreter versions and a Travis-CI timeout of 50 # minutes for a shard. PY27 = "2.7.15" PY35 = "3.5.6" PY36 = "3.6.6" _VERSIONS = (PY27, PY35, PY36) # This is the filename of a sentinel file that sits in the pyenv root directory. # Its purpose is to indicate whether pyenv has the correct interpreters installed # and will be useful for indicating whether we should trigger a reclone to update # pyenv. _INTERPRETER_SET_FINGERPRINT = "_".join(_VERSIONS) + "_pex_fingerprint" def ensure_python_distribution(version): # type: (str) -> Tuple[str, str, Callable[[Iterable[str]], Text]] if version not in _VERSIONS: raise ValueError("Please constrain version to one of {}".format(_VERSIONS)) pyenv_root = os.path.join(os.getcwd(), ".pyenv_test") interpreter_location = os.path.join(pyenv_root, "versions", version) pyenv = os.path.join(pyenv_root, "bin", "pyenv") pyenv_env = os.environ.copy() pyenv_env["PYENV_ROOT"] = pyenv_root pip = os.path.join(interpreter_location, "bin", "pip") if not os.path.exists(os.path.join(pyenv_root, _INTERPRETER_SET_FINGERPRINT)): bootstrap_python_installer(pyenv_root) if not os.path.exists(interpreter_location): env = pyenv_env.copy() if sys.platform.lower() == "linux": env["CONFIGURE_OPTS"] = "--enable-shared" subprocess.check_call([pyenv, "install", "--keep", version], env=env) subprocess.check_call([pip, "install", "-U", "pip"]) python = os.path.join(interpreter_location, "bin", "python" + version[0:3]) def run_pyenv(args): # type: (Iterable[str]) -> Text return to_unicode(subprocess.check_output([pyenv] + list(args), env=pyenv_env)) return python, pip, run_pyenv def ensure_python_interpreter(version): # type: (str) -> str python, _, _ = ensure_python_distribution(version) return python @contextmanager def environment_as(**kwargs): # type: (**str) -> Iterator[None] existing = {key: os.environ.get(key) for key in kwargs} def adjust_environment(mapping): for key, value in mapping.items(): if value is not None: os.environ[key] = value else: del os.environ[key] adjust_environment(kwargs) try: yield finally: adjust_environment(existing)