test_pkgutil.py revision 1e88992224d15bf3dbc30f058923b9f4a9d7149a 
1from test.test_support import run_unittest
2import unittest
3import sys
4import imp
5import pkgutil
6import os
7import os.path
8import tempfile
9import shutil
10import zipfile
11
12
13
14class 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    def test_getdata_zipfile(self):
52        zip = 'test_getdata_zipfile.zip'
53        pkg = 'test_getdata_zipfile'
54
55        # Include a LF and a CRLF, to test that binary data is read back
56        RESOURCE_DATA = 'Hello, world!\nSecond line\r\nThird line'
57
58        # Make a package with some resources
59        zip_file = os.path.join(self.dirname, zip)
60        z = zipfile.ZipFile(zip_file, 'w')
61
62        # Empty init.py
63        z.writestr(pkg + '/__init__.py', "")
64        # Resource files, res.txt, sub/res.txt
65        z.writestr(pkg + '/res.txt', RESOURCE_DATA)
66        z.writestr(pkg + '/sub/res.txt', RESOURCE_DATA)
67        z.close()
68
69        # Check we can read the resources
70        sys.path.insert(0, zip_file)
71        res1 = pkgutil.get_data(pkg, 'res.txt')
72        self.assertEqual(res1, RESOURCE_DATA)
73        res2 = pkgutil.get_data(pkg, 'sub/res.txt')
74        self.assertEqual(res2, RESOURCE_DATA)
75        del sys.path[0]
76
77class PkgutilPEP302Tests(unittest.TestCase):
78
79    class MyTestLoader(object):
80        def load_module(self, fullname):
81            # Create an empty module
82            mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
83            mod.__file__ = "<%s>" % self.__class__.__name__
84            mod.__loader__ = self
85            # Make it a package
86            mod.__path__ = []
87            # Count how many times the module is reloaded
88            mod.__dict__['loads'] = mod.__dict__.get('loads',0) + 1
89            return mod
90
91        def get_data(self, path):
92            return "Hello, world!"
93
94    class MyTestImporter(object):
95        def find_module(self, fullname, path=None):
96            return PkgutilPEP302Tests.MyTestLoader()
97
98    def setUp(self):
99        sys.meta_path.insert(0, self.MyTestImporter())
100
101    def tearDown(self):
102        del sys.meta_path[0]
103
104    def test_getdata_pep302(self):
105        # Use a dummy importer/loader
106        self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
107        del sys.modules['foo']
108
109    def test_alreadyloaded(self):
110        # Ensure that get_data works without reloading - the "loads" module
111        # variable in the example loader should count how many times a reload
112        # occurs.
113        import foo
114        self.assertEqual(foo.loads, 1)
115        self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
116        self.assertEqual(foo.loads, 1)
117        del sys.modules['foo']
118
119def test_main():
120    run_unittest(PkgutilTests, PkgutilPEP302Tests)
121
122if __name__ == '__main__':
123    test_main()
124