Molsen/add pex root option

pex/bin/pex.py

@@ -503,11 +503,18 @@ def make_relative_to_root(path):
   return os.path.normpath(path.format(pex_root=ENV.PEX_ROOT))
 
 
-def main():
+def main(args=None, log_callback=None):
+  if log_callback is not None:
+    log = log_callback
+
   parser, resolver_options_builder = configure_clp()
 
-  # split arguments early because optparse is dumb
-  args = sys.argv[1:]
+  if args is None:
+    # split arguments early because optparse is dumb
+    args = sys.argv[1:]
+  else:
+    args = args[:]
+
   try:
     separator = args.index('--')
     args, cmdline = args[:separator], args[separator + 1:]
@@ -542,7 +549,7 @@ def main():
 
     log('Running PEX file at %s with args %s' % (pex_builder.path(), cmdline), v=options.verbosity)
     pex = PEX(pex_builder.path(), interpreter=pex_builder.interpreter)
-    sys.exit(pex.run(args=list(cmdline)))
+    raise SystemExit(pex.run(args=list(cmdline)))
 
 
 if __name__ == '__main__':

pex/testing.py

@@ -8,8 +8,10 @@
 import sys
 import tempfile
 import zipfile
+from collections import namedtuple
 from textwrap import dedent
 
+from .bin.pex import main
 from .common import safe_mkdir, safe_rmtree
 from .compatibility import nested
 from .installer import EggInstaller, Packager
@@ -173,6 +175,42 @@ def write_simple_pex(td, exe_contents, dists=None, coverage=False):
   return pb
 
 
+class IntegResults(namedtuple('results', 'output return_code exception')):
+  """Convenience object to return integration run results."""
+
+  def assert_success(self):
+    assert self.exception is None and self.return_code is None
+
+  def assert_failure(self):
+    assert self.exception or self.return_code
+
+
+def run_pex_command(args, env=None):
+  """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.
+  """
+  def logger_callback(_output):
+    def mock_logger(msg, v=None):
+      _output.append(msg)
+
+    return mock_logger
+
+  exception = None
+  error_code = None
+  output = []
+
+  try:
+    main(args=args, log_callback=logger_callback(output))
+  except SystemExit as e:
+    error_code = e.code
+  except Exception as e:
+    exception = e
+  return IntegResults(output, error_code, exception)
+
+
 # TODO(wickman) Why not PEX.run?
 def run_simple_pex(pex, args=(), env=None):
   po = subprocess.Popen(

tests/test_integration.py

@@ -8,7 +8,7 @@
 from twitter.common.contextutil import environment_as, temporary_dir
 
 from pex.compatibility import WINDOWS
-from pex.testing import run_simple_pex_test
+from pex.testing import run_pex_command, run_simple_pex_test
 from pex.util import named_temporary_file
 
 
@@ -23,6 +23,23 @@ def test_pex_raise():
   run_simple_pex_test(body, coverage=True)
 
 
+def test_pex_root():
+  with temporary_dir() as tmp_home:
+    with environment_as(HOME=tmp_home):
+      with temporary_dir() as td:
+        with temporary_dir() as output_dir:
+          env = os.environ.copy()
+          env['PEX_INTERPRETER'] = '1'
+
+          output_path = os.path.join(output_dir, 'pex.pex')
+          args = ['pex', '-o', output_path, '--not-zip-safe', '--pex-root={0}'.format(td)]
+          results = run_pex_command(args=args, env=env)
+          results.assert_success()
+          assert ['pex.pex'] == os.listdir(output_dir), 'Expected built pex file.'
+          assert [] == os.listdir(tmp_home), 'Expected empty temp home dir.'
+          assert 'build' in os.listdir(td), 'Expected build directory in tmp pex root.'
+
+
 def test_pex_interpreter():
   with named_temporary_file() as fp:
     fp.write(b"print('Hello world')")