Lib/test/test_pkgutil.py

   1 from test.test_support import run_unittest
   2 import unittest
   3 import sys
   4 import imp
   5 import pkgutil
   6 import os
   7 import os.path
   8 import tempfile
   9 import shutil
  10 import zipfile
  11
  12
  13
  14 class PkgutilTests(unittest.TestCase):
  15
  16     def setUp(self):
  17         self.dirname = tempfile.mkdtemp()
  18         sys.path.insert(0, self.dirname)
  19
  20     def tearDown(self):
  21         del sys.path[0]
  22         shutil.rmtree(self.dirname)
  23
  24     def test_getdata_filesys(self):
  25         pkg = 'test_getdata_filesys'
  26
  27         # Include a LF and a CRLF, to test that binary data is read back
  28         RESOURCE_DATA = 'Hello, world!\nSecond line\r\nThird line'
  29
  30         # Make a package with some resources
  31         package_dir = os.path.join(self.dirname, pkg)
  32         os.mkdir(package_dir)
  33         # Empty init.py
  34         f = open(os.path.join(package_dir, '__init__.py'), "wb")
  35         f.close()
  36         # Resource files, res.txt, sub/res.txt
  37         f = open(os.path.join(package_dir, 'res.txt'), "wb")
  38         f.write(RESOURCE_DATA)
  39         f.close()
  40         os.mkdir(os.path.join(package_dir, 'sub'))
  41         f = open(os.path.join(package_dir, 'sub', 'res.txt'), "wb")
  42         f.write(RESOURCE_DATA)
  43         f.close()
  44
  45         # Check we can read the resources
  46         res1 = pkgutil.get_data(pkg, 'res.txt')
  47         self.assertEqual(res1, RESOURCE_DATA)
  48         res2 = pkgutil.get_data(pkg, 'sub/res.txt')
  49         self.assertEqual(res2, RESOURCE_DATA)
  50
  51         del sys.modules[pkg]
  52
  53     def test_getdata_zipfile(self):
  54         zip = 'test_getdata_zipfile.zip'
  55         pkg = 'test_getdata_zipfile'
  56
  57         # Include a LF and a CRLF, to test that binary data is read back
  58         RESOURCE_DATA = 'Hello, world!\nSecond line\r\nThird line'
  59
  60         # Make a package with some resources
  61         zip_file = os.path.join(self.dirname, zip)
  62         z = zipfile.ZipFile(zip_file, 'w')
  63
  64         # Empty init.py
  65         z.writestr(pkg + '/__init__.py', "")
  66         # Resource files, res.txt, sub/res.txt
  67         z.writestr(pkg + '/res.txt', RESOURCE_DATA)
  68         z.writestr(pkg + '/sub/res.txt', RESOURCE_DATA)
  69         z.close()
  70
  71         # Check we can read the resources
  72         sys.path.insert(0, zip_file)
  73         res1 = pkgutil.get_data(pkg, 'res.txt')
  74         self.assertEqual(res1, RESOURCE_DATA)
  75         res2 = pkgutil.get_data(pkg, 'sub/res.txt')
  76         self.assertEqual(res2, RESOURCE_DATA)
  77         del sys.path[0]
  78
  79         del sys.modules[pkg]
  80
  81 class PkgutilPEP302Tests(unittest.TestCase):
  82
  83     class MyTestLoader(object):
  84         def load_module(self, fullname):
  85             # Create an empty module
  86             mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
  87             mod.__file__ = "<%s>" % self.__class__.__name__
  88             mod.__loader__ = self
  89             # Make it a package
  90             mod.__path__ = []
  91             # Count how many times the module is reloaded
  92             mod.__dict__['loads'] = mod.__dict__.get('loads',0) + 1
  93             return mod
  94
  95         def get_data(self, path):
  96             return "Hello, world!"
  97
  98     class MyTestImporter(object):
  99         def find_module(self, fullname, path=None):
 100             return PkgutilPEP302Tests.MyTestLoader()
 101
 102     def setUp(self):
 103         sys.meta_path.insert(0, self.MyTestImporter())
 104
 105     def tearDown(self):
 106         del sys.meta_path[0]
 107
 108     def test_getdata_pep302(self):
 109         # Use a dummy importer/loader
 110         self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
 111         del sys.modules['foo']
 112
 113     def test_alreadyloaded(self):
 114         # Ensure that get_data works without reloading - the "loads" module
 115         # variable in the example loader should count how many times a reload
 116         # occurs.
 117         import foo
 118         self.assertEqual(foo.loads, 1)
 119         self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
 120         self.assertEqual(foo.loads, 1)
 121         del sys.modules['foo']
 122
 123 def test_main():
 124     run_unittest(PkgutilTests, PkgutilPEP302Tests)
 125     # this is necessary if test is run repeated (like when finding leaks)
 126     import zipimport
 127     zipimport._zip_directory_cache.clear()
 128
 129 if __name__ == '__main__':
 130     test_main()