Lib/test/script_helper.py

   1 # Common utility functions used by various script execution tests
   2 #  e.g. test_cmd_line, test_cmd_line_script and test_runpy
   3
   4 import sys
   5 import os
   6 import os.path
   7 import tempfile
   8 import subprocess
   9 import py_compile
  10 import contextlib
  11 import shutil
  12 import zipfile
  13
  14 # Executing the interpreter in a subprocess
  15 def _assert_python(expected_success, *args, **env_vars):
  16     cmd_line = [sys.executable]
  17     if not env_vars:
  18         cmd_line.append('-E')
  19     cmd_line.extend(args)
  20     # Need to preserve the original environment, for in-place testing of
  21     # shared library builds.
  22     env = os.environ.copy()
  23     env.update(env_vars)
  24     p = subprocess.Popen(cmd_line, stdin=subprocess.PIPE,
  25                          stdout=subprocess.PIPE, stderr=subprocess.PIPE,
  26                          env=env)
  27     try:
  28         out, err = p.communicate()
  29     finally:
  30         subprocess._cleanup()
  31         p.stdout.close()
  32         p.stderr.close()
  33     rc = p.returncode
  34     if (rc and expected_success) or (not rc and not expected_success):
  35         raise AssertionError(
  36             "Process return code is %d, "
  37             "stderr follows:\n%s" % (rc, err.decode('ascii', 'ignore')))
  38     return rc, out, err
  39
  40 def assert_python_ok(*args, **env_vars):
  41     """
  42     Assert that running the interpreter with `args` and optional environment
  43     variables `env_vars` is ok and return a (return code, stdout, stderr) tuple.
  44     """
  45     return _assert_python(True, *args, **env_vars)
  46
  47 def assert_python_failure(*args, **env_vars):
  48     """
  49     Assert that running the interpreter with `args` and optional environment
  50     variables `env_vars` fails and return a (return code, stdout, stderr) tuple.
  51     """
  52     return _assert_python(False, *args, **env_vars)
  53
  54 def python_exit_code(*args):
  55     cmd_line = [sys.executable, '-E']
  56     cmd_line.extend(args)
  57     with open(os.devnull, 'w') as devnull:
  58         return subprocess.call(cmd_line, stdout=devnull,
  59                                 stderr=subprocess.STDOUT)
  60
  61 def spawn_python(*args, **kwargs):
  62     cmd_line = [sys.executable, '-E']
  63     cmd_line.extend(args)
  64     return subprocess.Popen(cmd_line, stdin=subprocess.PIPE,
  65                             stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
  66                             **kwargs)
  67
  68 def kill_python(p):
  69     p.stdin.close()
  70     data = p.stdout.read()
  71     p.stdout.close()
  72     # try to cleanup the child so we don't appear to leak when running
  73     # with regrtest -R.
  74     p.wait()
  75     subprocess._cleanup()
  76     return data
  77
  78 def run_python(*args, **kwargs):
  79     if __debug__:
  80         p = spawn_python(*args, **kwargs)
  81     else:
  82         p = spawn_python('-O', *args, **kwargs)
  83     stdout_data = kill_python(p)
  84     return p.wait(), stdout_data
  85
  86 # Script creation utilities
  87 @contextlib.contextmanager
  88 def temp_dir():
  89     dirname = tempfile.mkdtemp()
  90     dirname = os.path.realpath(dirname)
  91     try:
  92         yield dirname
  93     finally:
  94         shutil.rmtree(dirname)
  95
  96 def make_script(script_dir, script_basename, source):
  97     script_filename = script_basename+os.extsep+'py'
  98     script_name = os.path.join(script_dir, script_filename)
  99     script_file = open(script_name, 'w')
 100     script_file.write(source)
 101     script_file.close()
 102     return script_name
 103
 104 def compile_script(script_name):
 105     py_compile.compile(script_name, doraise=True)
 106     if __debug__:
 107         compiled_name = script_name + 'c'
 108     else:
 109         compiled_name = script_name + 'o'
 110     return compiled_name
 111
 112 def make_zip_script(zip_dir, zip_basename, script_name, name_in_zip=None):
 113     zip_filename = zip_basename+os.extsep+'zip'
 114     zip_name = os.path.join(zip_dir, zip_filename)
 115     zip_file = zipfile.ZipFile(zip_name, 'w')
 116     if name_in_zip is None:
 117         name_in_zip = os.path.basename(script_name)
 118     zip_file.write(script_name, name_in_zip)
 119     zip_file.close()
 120     #if test.test_support.verbose:
 121     #    zip_file = zipfile.ZipFile(zip_name, 'r')
 122     #    print 'Contents of %r:' % zip_name
 123     #    zip_file.printdir()
 124     #    zip_file.close()
 125     return zip_name, os.path.join(zip_name, name_in_zip)
 126
 127 def make_pkg(pkg_dir):
 128     os.mkdir(pkg_dir)
 129     make_script(pkg_dir, '__init__', '')
 130
 131 def make_zip_pkg(zip_dir, zip_basename, pkg_name, script_basename,
 132                  source, depth=1, compiled=False):
 133     unlink = []
 134     init_name = make_script(zip_dir, '__init__', '')
 135     unlink.append(init_name)
 136     init_basename = os.path.basename(init_name)
 137     script_name = make_script(zip_dir, script_basename, source)
 138     unlink.append(script_name)
 139     if compiled:
 140         init_name = compile_script(init_name)
 141         script_name = compile_script(script_name)
 142         unlink.extend((init_name, script_name))
 143     pkg_names = [os.sep.join([pkg_name]*i) for i in range(1, depth+1)]
 144     script_name_in_zip = os.path.join(pkg_names[-1], os.path.basename(script_name))
 145     zip_filename = zip_basename+os.extsep+'zip'
 146     zip_name = os.path.join(zip_dir, zip_filename)
 147     zip_file = zipfile.ZipFile(zip_name, 'w')
 148     for name in pkg_names:
 149         init_name_in_zip = os.path.join(name, init_basename)
 150         zip_file.write(init_name, init_name_in_zip)
 151     zip_file.write(script_name, script_name_in_zip)
 152     zip_file.close()
 153     for name in unlink:
 154         os.unlink(name)
 155     #if test.test_support.verbose:
 156     #    zip_file = zipfile.ZipFile(zip_name, 'r')
 157     #    print 'Contents of %r:' % zip_name
 158     #    zip_file.printdir()
 159     #    zip_file.close()
 160     return zip_name, os.path.join(zip_name, script_name_in_zip)