Lib/test/test_cmd_line_script.py

   1 # Tests command line execution of scripts
   2
   3 import unittest
   4 import os
   5 import os.path
   6 import sys
   7 import test.test_support
   8 import tempfile
   9 import subprocess
  10 import py_compile
  11 import contextlib
  12 import shutil
  13 import zipfile
  14
  15 verbose = test.test_support.verbose
  16
  17 # XXX ncoghlan: Should we consider moving these to test_support?
  18 from test_cmd_line import _spawn_python, _kill_python
  19
  20 def _run_python(*args):
  21     if __debug__:
  22         p = _spawn_python(*args)
  23     else:
  24         p = _spawn_python('-O', *args)
  25     stdout_data = _kill_python(p)
  26     return p.wait(), stdout_data
  27
  28 @contextlib.contextmanager
  29 def temp_dir():
  30     dirname = tempfile.mkdtemp()
  31     dirname = os.path.realpath(dirname)
  32     try:
  33         yield dirname
  34     finally:
  35         shutil.rmtree(dirname)
  36
  37 test_source = """\
  38 # Script may be run with optimisation enabled, so don't rely on assert
  39 # statements being executed
  40 def assertEqual(lhs, rhs):
  41     if lhs != rhs:
  42         raise AssertionError('%r != %r' % (lhs, rhs))
  43 def assertIdentical(lhs, rhs):
  44     if lhs is not rhs:
  45         raise AssertionError('%r is not %r' % (lhs, rhs))
  46 # Check basic code execution
  47 result = ['Top level assignment']
  48 def f():
  49     result.append('Lower level reference')
  50 f()
  51 assertEqual(result, ['Top level assignment', 'Lower level reference'])
  52 # Check population of magic variables
  53 assertEqual(__name__, '__main__')
  54 print '__file__==%r' % __file__
  55 print '__package__==%r' % __package__
  56 # Check the sys module
  57 import sys
  58 assertIdentical(globals(), sys.modules[__name__].__dict__)
  59 print 'sys.argv[0]==%r' % sys.argv[0]
  60 """
  61
  62 def _make_test_script(script_dir, script_basename, source=test_source):
  63     script_filename = script_basename+os.extsep+'py'
  64     script_name = os.path.join(script_dir, script_filename)
  65     script_file = open(script_name, 'w')
  66     script_file.write(source)
  67     script_file.close()
  68     return script_name
  69
  70 def _compile_test_script(script_name):
  71     py_compile.compile(script_name, doraise=True)
  72     if __debug__:
  73         compiled_name = script_name + 'c'
  74     else:
  75         compiled_name = script_name + 'o'
  76     return compiled_name
  77
  78 def _make_test_zip(zip_dir, zip_basename, script_name, name_in_zip=None):
  79     zip_filename = zip_basename+os.extsep+'zip'
  80     zip_name = os.path.join(zip_dir, zip_filename)
  81     zip_file = zipfile.ZipFile(zip_name, 'w')
  82     if name_in_zip is None:
  83         name_in_zip = os.path.basename(script_name)
  84     zip_file.write(script_name, name_in_zip)
  85     zip_file.close()
  86     #if verbose:
  87     #    zip_file = zipfile.ZipFile(zip_name, 'r')
  88     #    print 'Contents of %r:' % zip_name
  89     #    zip_file.printdir()
  90     #    zip_file.close()
  91     return zip_name, os.path.join(zip_name, name_in_zip)
  92
  93 def _make_test_pkg(pkg_dir):
  94     os.mkdir(pkg_dir)
  95     _make_test_script(pkg_dir, '__init__', '')
  96
  97 def _make_test_zip_pkg(zip_dir, zip_basename, pkg_name, script_basename,
  98                        source=test_source, depth=1):
  99     init_name = _make_test_script(zip_dir, '__init__', '')
 100     init_basename = os.path.basename(init_name)
 101     script_name = _make_test_script(zip_dir, script_basename, source)
 102     pkg_names = [os.sep.join([pkg_name]*i) for i in range(1, depth+1)]
 103     script_name_in_zip = os.path.join(pkg_names[-1], os.path.basename(script_name))
 104     zip_filename = zip_basename+os.extsep+'zip'
 105     zip_name = os.path.join(zip_dir, zip_filename)
 106     zip_file = zipfile.ZipFile(zip_name, 'w')
 107     for name in pkg_names:
 108         init_name_in_zip = os.path.join(name, init_basename)
 109         zip_file.write(init_name, init_name_in_zip)
 110     zip_file.write(script_name, script_name_in_zip)
 111     zip_file.close()
 112     os.unlink(init_name)
 113     os.unlink(script_name)
 114     #if verbose:
 115     #    zip_file = zipfile.ZipFile(zip_name, 'r')
 116     #    print 'Contents of %r:' % zip_name
 117     #    zip_file.printdir()
 118     #    zip_file.close()
 119     return zip_name, os.path.join(zip_name, script_name_in_zip)
 120
 121 # There's no easy way to pass the script directory in to get
 122 # -m to work (avoiding that is the whole point of making
 123 # directories and zipfiles executable!)
 124 # So we fake it for testing purposes with a custom launch script
 125 launch_source = """\
 126 import sys, os.path, runpy
 127 sys.path.insert(0, %s)
 128 runpy._run_module_as_main(%r)
 129 """
 130
 131 def _make_launch_script(script_dir, script_basename, module_name, path=None):
 132     if path is None:
 133         path = "os.path.dirname(__file__)"
 134     else:
 135         path = repr(path)
 136     source = launch_source % (path, module_name)
 137     return _make_test_script(script_dir, script_basename, source)
 138
 139 class CmdLineTest(unittest.TestCase):
 140     def _check_script(self, script_name, expected_file,
 141                             expected_argv0, expected_package,
 142                             *cmd_line_switches):
 143         run_args = cmd_line_switches + (script_name,)
 144         exit_code, data = _run_python(*run_args)
 145         if verbose:
 146             print 'Output from test script %r:' % script_name
 147             print data
 148         self.assertEqual(exit_code, 0)
 149         printed_file = '__file__==%r' % expected_file
 150         printed_argv0 = 'sys.argv[0]==%r' % expected_argv0
 151         printed_package = '__package__==%r' % expected_package
 152         if verbose:
 153             print 'Expected output:'
 154             print printed_file
 155             print printed_package
 156             print printed_argv0
 157         self.assertTrue(printed_file in data)
 158         self.assertTrue(printed_package in data)
 159         self.assertTrue(printed_argv0 in data)
 160
 161     def _check_import_error(self, script_name, expected_msg,
 162                             *cmd_line_switches):
 163         run_args = cmd_line_switches + (script_name,)
 164         exit_code, data = _run_python(*run_args)
 165         if verbose:
 166             print 'Output from test script %r:' % script_name
 167             print data
 168             print 'Expected output: %r' % expected_msg
 169         self.assertTrue(expected_msg in data)
 170
 171     def test_basic_script(self):
 172         with temp_dir() as script_dir:
 173             script_name = _make_test_script(script_dir, 'script')
 174             self._check_script(script_name, script_name, script_name, None)
 175
 176     def test_script_compiled(self):
 177         with temp_dir() as script_dir:
 178             script_name = _make_test_script(script_dir, 'script')
 179             compiled_name = _compile_test_script(script_name)
 180             os.remove(script_name)
 181             self._check_script(compiled_name, compiled_name, compiled_name, None)
 182
 183     def test_directory(self):
 184         with temp_dir() as script_dir:
 185             script_name = _make_test_script(script_dir, '__main__')
 186             self._check_script(script_dir, script_name, script_dir, '')
 187
 188     def test_directory_compiled(self):
 189         with temp_dir() as script_dir:
 190             script_name = _make_test_script(script_dir, '__main__')
 191             compiled_name = _compile_test_script(script_name)
 192             os.remove(script_name)
 193             self._check_script(script_dir, compiled_name, script_dir, '')
 194
 195     def test_directory_error(self):
 196         with temp_dir() as script_dir:
 197             msg = "can't find '__main__.py' in %r" % script_dir
 198             self._check_import_error(script_dir, msg)
 199
 200     def test_zipfile(self):
 201         with temp_dir() as script_dir:
 202             script_name = _make_test_script(script_dir, '__main__')
 203             zip_name, run_name = _make_test_zip(script_dir, 'test_zip', script_name)
 204             self._check_script(zip_name, run_name, zip_name, '')
 205
 206     def test_zipfile_compiled(self):
 207         with temp_dir() as script_dir:
 208             script_name = _make_test_script(script_dir, '__main__')
 209             compiled_name = _compile_test_script(script_name)
 210             zip_name, run_name = _make_test_zip(script_dir, 'test_zip', compiled_name)
 211             self._check_script(zip_name, run_name, zip_name, '')
 212
 213     def test_zipfile_error(self):
 214         with temp_dir() as script_dir:
 215             script_name = _make_test_script(script_dir, 'not_main')
 216             zip_name, run_name = _make_test_zip(script_dir, 'test_zip', script_name)
 217             msg = "can't find '__main__.py' in %r" % zip_name
 218             self._check_import_error(zip_name, msg)
 219
 220     def test_module_in_package(self):
 221         with temp_dir() as script_dir:
 222             pkg_dir = os.path.join(script_dir, 'test_pkg')
 223             _make_test_pkg(pkg_dir)
 224             script_name = _make_test_script(pkg_dir, 'script')
 225             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg.script')
 226             self._check_script(launch_name, script_name, script_name, 'test_pkg')
 227
 228     def test_module_in_package_in_zipfile(self):
 229         with temp_dir() as script_dir:
 230             zip_name, run_name = _make_test_zip_pkg(script_dir, 'test_zip', 'test_pkg', 'script')
 231             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg.script', zip_name)
 232             self._check_script(launch_name, run_name, run_name, 'test_pkg')
 233
 234     def test_module_in_subpackage_in_zipfile(self):
 235         with temp_dir() as script_dir:
 236             zip_name, run_name = _make_test_zip_pkg(script_dir, 'test_zip', 'test_pkg', 'script', depth=2)
 237             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg.test_pkg.script', zip_name)
 238             self._check_script(launch_name, run_name, run_name, 'test_pkg.test_pkg')
 239
 240     def test_package(self):
 241         with temp_dir() as script_dir:
 242             pkg_dir = os.path.join(script_dir, 'test_pkg')
 243             _make_test_pkg(pkg_dir)
 244             script_name = _make_test_script(pkg_dir, '__main__')
 245             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg')
 246             self._check_script(launch_name, script_name,
 247                                script_name, 'test_pkg')
 248
 249     def test_package_compiled(self):
 250         with temp_dir() as script_dir:
 251             pkg_dir = os.path.join(script_dir, 'test_pkg')
 252             _make_test_pkg(pkg_dir)
 253             script_name = _make_test_script(pkg_dir, '__main__')
 254             compiled_name = _compile_test_script(script_name)
 255             os.remove(script_name)
 256             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg')
 257             self._check_script(launch_name, compiled_name,
 258                                compiled_name, 'test_pkg')
 259
 260     def test_package_error(self):
 261         with temp_dir() as script_dir:
 262             pkg_dir = os.path.join(script_dir, 'test_pkg')
 263             _make_test_pkg(pkg_dir)
 264             msg = ("'test_pkg' is a package and cannot "
 265                    "be directly executed")
 266             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg')
 267             self._check_import_error(launch_name, msg)
 268
 269     def test_package_recursion(self):
 270         with temp_dir() as script_dir:
 271             pkg_dir = os.path.join(script_dir, 'test_pkg')
 272             _make_test_pkg(pkg_dir)
 273             main_dir = os.path.join(pkg_dir, '__main__')
 274             _make_test_pkg(main_dir)
 275             msg = ("Cannot use package as __main__ module; "
 276                    "'test_pkg' is a package and cannot "
 277                    "be directly executed")
 278             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg')
 279             self._check_import_error(launch_name, msg)
 280
 281
 282 def test_main():
 283     test.test_support.run_unittest(CmdLineTest)
 284     test.test_support.reap_children()
 285
 286 if __name__ == '__main__':
 287     test_main()