1"""Test harness for the zipapp module."""
2
3import io
4import pathlib
5import stat
6import sys
7import tempfile
8import unittest
9import zipapp
10import zipfile
11
12from unittest.mock import patch
13
14class ZipAppTest(unittest.TestCase):
15
16    """Test zipapp module functionality."""
17
18    def setUp(self):
19        tmpdir = tempfile.TemporaryDirectory()
20        self.addCleanup(tmpdir.cleanup)
21        self.tmpdir = pathlib.Path(tmpdir.name)
22
23    def test_create_archive(self):
24        # Test packing a directory.
25        source = self.tmpdir / 'source'
26        source.mkdir()
27        (source / '__main__.py').touch()
28        target = self.tmpdir / 'source.pyz'
29        zipapp.create_archive(str(source), str(target))
30        self.assertTrue(target.is_file())
31
32    def test_create_archive_with_pathlib(self):
33        # Test packing a directory using Path objects for source and target.
34        source = self.tmpdir / 'source'
35        source.mkdir()
36        (source / '__main__.py').touch()
37        target = self.tmpdir / 'source.pyz'
38        zipapp.create_archive(source, target)
39        self.assertTrue(target.is_file())
40
41    def test_create_archive_with_subdirs(self):
42        # Test packing a directory includes entries for subdirectories.
43        source = self.tmpdir / 'source'
44        source.mkdir()
45        (source / '__main__.py').touch()
46        (source / 'foo').mkdir()
47        (source / 'bar').mkdir()
48        (source / 'foo' / '__init__.py').touch()
49        target = io.BytesIO()
50        zipapp.create_archive(str(source), target)
51        target.seek(0)
52        with zipfile.ZipFile(target, 'r') as z:
53            self.assertIn('foo/', z.namelist())
54            self.assertIn('bar/', z.namelist())
55
56    def test_create_archive_default_target(self):
57        # Test packing a directory to the default name.
58        source = self.tmpdir / 'source'
59        source.mkdir()
60        (source / '__main__.py').touch()
61        zipapp.create_archive(str(source))
62        expected_target = self.tmpdir / 'source.pyz'
63        self.assertTrue(expected_target.is_file())
64
65    def test_no_main(self):
66        # Test that packing a directory with no __main__.py fails.
67        source = self.tmpdir / 'source'
68        source.mkdir()
69        (source / 'foo.py').touch()
70        target = self.tmpdir / 'source.pyz'
71        with self.assertRaises(zipapp.ZipAppError):
72            zipapp.create_archive(str(source), str(target))
73
74    def test_main_and_main_py(self):
75        # Test that supplying a main argument with __main__.py fails.
76        source = self.tmpdir / 'source'
77        source.mkdir()
78        (source / '__main__.py').touch()
79        target = self.tmpdir / 'source.pyz'
80        with self.assertRaises(zipapp.ZipAppError):
81            zipapp.create_archive(str(source), str(target), main='pkg.mod:fn')
82
83    def test_main_written(self):
84        # Test that the __main__.py is written correctly.
85        source = self.tmpdir / 'source'
86        source.mkdir()
87        (source / 'foo.py').touch()
88        target = self.tmpdir / 'source.pyz'
89        zipapp.create_archive(str(source), str(target), main='pkg.mod:fn')
90        with zipfile.ZipFile(str(target), 'r') as z:
91            self.assertIn('__main__.py', z.namelist())
92            self.assertIn(b'pkg.mod.fn()', z.read('__main__.py'))
93
94    def test_main_only_written_once(self):
95        # Test that we don't write multiple __main__.py files.
96        # The initial implementation had this bug; zip files allow
97        # multiple entries with the same name
98        source = self.tmpdir / 'source'
99        source.mkdir()
100        # Write 2 files, as the original bug wrote __main__.py
101        # once for each file written :-(
102        # See http://bugs.python.org/review/23491/diff/13982/Lib/zipapp.py#newcode67Lib/zipapp.py:67
103        # (line 67)
104        (source / 'foo.py').touch()
105        (source / 'bar.py').touch()
106        target = self.tmpdir / 'source.pyz'
107        zipapp.create_archive(str(source), str(target), main='pkg.mod:fn')
108        with zipfile.ZipFile(str(target), 'r') as z:
109            self.assertEqual(1, z.namelist().count('__main__.py'))
110
111    def test_main_validation(self):
112        # Test that invalid values for main are rejected.
113        source = self.tmpdir / 'source'
114        source.mkdir()
115        target = self.tmpdir / 'source.pyz'
116        problems = [
117            '', 'foo', 'foo:', ':bar', '12:bar', 'a.b.c.:d',
118            '.a:b', 'a:b.', 'a:.b', 'a:silly name'
119        ]
120        for main in problems:
121            with self.subTest(main=main):
122                with self.assertRaises(zipapp.ZipAppError):
123                    zipapp.create_archive(str(source), str(target), main=main)
124
125    def test_default_no_shebang(self):
126        # Test that no shebang line is written to the target by default.
127        source = self.tmpdir / 'source'
128        source.mkdir()
129        (source / '__main__.py').touch()
130        target = self.tmpdir / 'source.pyz'
131        zipapp.create_archive(str(source), str(target))
132        with target.open('rb') as f:
133            self.assertNotEqual(f.read(2), b'#!')
134
135    def test_custom_interpreter(self):
136        # Test that a shebang line with a custom interpreter is written
137        # correctly.
138        source = self.tmpdir / 'source'
139        source.mkdir()
140        (source / '__main__.py').touch()
141        target = self.tmpdir / 'source.pyz'
142        zipapp.create_archive(str(source), str(target), interpreter='python')
143        with target.open('rb') as f:
144            self.assertEqual(f.read(2), b'#!')
145            self.assertEqual(b'python\n', f.readline())
146
147    def test_pack_to_fileobj(self):
148        # Test that we can pack to a file object.
149        source = self.tmpdir / 'source'
150        source.mkdir()
151        (source / '__main__.py').touch()
152        target = io.BytesIO()
153        zipapp.create_archive(str(source), target, interpreter='python')
154        self.assertTrue(target.getvalue().startswith(b'#!python\n'))
155
156    def test_read_shebang(self):
157        # Test that we can read the shebang line correctly.
158        source = self.tmpdir / 'source'
159        source.mkdir()
160        (source / '__main__.py').touch()
161        target = self.tmpdir / 'source.pyz'
162        zipapp.create_archive(str(source), str(target), interpreter='python')
163        self.assertEqual(zipapp.get_interpreter(str(target)), 'python')
164
165    def test_read_missing_shebang(self):
166        # Test that reading the shebang line of a file without one returns None.
167        source = self.tmpdir / 'source'
168        source.mkdir()
169        (source / '__main__.py').touch()
170        target = self.tmpdir / 'source.pyz'
171        zipapp.create_archive(str(source), str(target))
172        self.assertEqual(zipapp.get_interpreter(str(target)), None)
173
174    def test_modify_shebang(self):
175        # Test that we can change the shebang of a file.
176        source = self.tmpdir / 'source'
177        source.mkdir()
178        (source / '__main__.py').touch()
179        target = self.tmpdir / 'source.pyz'
180        zipapp.create_archive(str(source), str(target), interpreter='python')
181        new_target = self.tmpdir / 'changed.pyz'
182        zipapp.create_archive(str(target), str(new_target), interpreter='python2.7')
183        self.assertEqual(zipapp.get_interpreter(str(new_target)), 'python2.7')
184
185    def test_write_shebang_to_fileobj(self):
186        # Test that we can change the shebang of a file, writing the result to a
187        # file object.
188        source = self.tmpdir / 'source'
189        source.mkdir()
190        (source / '__main__.py').touch()
191        target = self.tmpdir / 'source.pyz'
192        zipapp.create_archive(str(source), str(target), interpreter='python')
193        new_target = io.BytesIO()
194        zipapp.create_archive(str(target), new_target, interpreter='python2.7')
195        self.assertTrue(new_target.getvalue().startswith(b'#!python2.7\n'))
196
197    def test_read_from_pathobj(self):
198        # Test that we can copy an archive using a pathlib.Path object
199        # for the source.
200        source = self.tmpdir / 'source'
201        source.mkdir()
202        (source / '__main__.py').touch()
203        target1 = self.tmpdir / 'target1.pyz'
204        target2 = self.tmpdir / 'target2.pyz'
205        zipapp.create_archive(source, target1, interpreter='python')
206        zipapp.create_archive(target1, target2, interpreter='python2.7')
207        self.assertEqual(zipapp.get_interpreter(target2), 'python2.7')
208
209    def test_read_from_fileobj(self):
210        # Test that we can copy an archive using an open file object.
211        source = self.tmpdir / 'source'
212        source.mkdir()
213        (source / '__main__.py').touch()
214        target = self.tmpdir / 'source.pyz'
215        temp_archive = io.BytesIO()
216        zipapp.create_archive(str(source), temp_archive, interpreter='python')
217        new_target = io.BytesIO()
218        temp_archive.seek(0)
219        zipapp.create_archive(temp_archive, new_target, interpreter='python2.7')
220        self.assertTrue(new_target.getvalue().startswith(b'#!python2.7\n'))
221
222    def test_remove_shebang(self):
223        # Test that we can remove the shebang from a file.
224        source = self.tmpdir / 'source'
225        source.mkdir()
226        (source / '__main__.py').touch()
227        target = self.tmpdir / 'source.pyz'
228        zipapp.create_archive(str(source), str(target), interpreter='python')
229        new_target = self.tmpdir / 'changed.pyz'
230        zipapp.create_archive(str(target), str(new_target), interpreter=None)
231        self.assertEqual(zipapp.get_interpreter(str(new_target)), None)
232
233    def test_content_of_copied_archive(self):
234        # Test that copying an archive doesn't corrupt it.
235        source = self.tmpdir / 'source'
236        source.mkdir()
237        (source / '__main__.py').touch()
238        target = io.BytesIO()
239        zipapp.create_archive(str(source), target, interpreter='python')
240        new_target = io.BytesIO()
241        target.seek(0)
242        zipapp.create_archive(target, new_target, interpreter=None)
243        new_target.seek(0)
244        with zipfile.ZipFile(new_target, 'r') as z:
245            self.assertEqual(set(z.namelist()), {'__main__.py'})
246
247    # (Unix only) tests that archives with shebang lines are made executable
248    @unittest.skipIf(sys.platform == 'win32',
249                     'Windows does not support an executable bit')
250    def test_shebang_is_executable(self):
251        # Test that an archive with a shebang line is made executable.
252        source = self.tmpdir / 'source'
253        source.mkdir()
254        (source / '__main__.py').touch()
255        target = self.tmpdir / 'source.pyz'
256        zipapp.create_archive(str(source), str(target), interpreter='python')
257        self.assertTrue(target.stat().st_mode & stat.S_IEXEC)
258
259    @unittest.skipIf(sys.platform == 'win32',
260                     'Windows does not support an executable bit')
261    def test_no_shebang_is_not_executable(self):
262        # Test that an archive with no shebang line is not made executable.
263        source = self.tmpdir / 'source'
264        source.mkdir()
265        (source / '__main__.py').touch()
266        target = self.tmpdir / 'source.pyz'
267        zipapp.create_archive(str(source), str(target), interpreter=None)
268        self.assertFalse(target.stat().st_mode & stat.S_IEXEC)
269
270
271class ZipAppCmdlineTest(unittest.TestCase):
272
273    """Test zipapp module command line API."""
274
275    def setUp(self):
276        tmpdir = tempfile.TemporaryDirectory()
277        self.addCleanup(tmpdir.cleanup)
278        self.tmpdir = pathlib.Path(tmpdir.name)
279
280    def make_archive(self):
281        # Test that an archive with no shebang line is not made executable.
282        source = self.tmpdir / 'source'
283        source.mkdir()
284        (source / '__main__.py').touch()
285        target = self.tmpdir / 'source.pyz'
286        zipapp.create_archive(source, target)
287        return target
288
289    def test_cmdline_create(self):
290        # Test the basic command line API.
291        source = self.tmpdir / 'source'
292        source.mkdir()
293        (source / '__main__.py').touch()
294        args = [str(source)]
295        zipapp.main(args)
296        target = source.with_suffix('.pyz')
297        self.assertTrue(target.is_file())
298
299    def test_cmdline_copy(self):
300        # Test copying an archive.
301        original = self.make_archive()
302        target = self.tmpdir / 'target.pyz'
303        args = [str(original), '-o', str(target)]
304        zipapp.main(args)
305        self.assertTrue(target.is_file())
306
307    def test_cmdline_copy_inplace(self):
308        # Test copying an archive in place fails.
309        original = self.make_archive()
310        target = self.tmpdir / 'target.pyz'
311        args = [str(original), '-o', str(original)]
312        with self.assertRaises(SystemExit) as cm:
313            zipapp.main(args)
314        # Program should exit with a non-zero returm code.
315        self.assertTrue(cm.exception.code)
316
317    def test_cmdline_copy_change_main(self):
318        # Test copying an archive doesn't allow changing __main__.py.
319        original = self.make_archive()
320        target = self.tmpdir / 'target.pyz'
321        args = [str(original), '-o', str(target), '-m', 'foo:bar']
322        with self.assertRaises(SystemExit) as cm:
323            zipapp.main(args)
324        # Program should exit with a non-zero returm code.
325        self.assertTrue(cm.exception.code)
326
327    @patch('sys.stdout', new_callable=io.StringIO)
328    def test_info_command(self, mock_stdout):
329        # Test the output of the info command.
330        target = self.make_archive()
331        args = [str(target), '--info']
332        with self.assertRaises(SystemExit) as cm:
333            zipapp.main(args)
334        # Program should exit with a zero returm code.
335        self.assertEqual(cm.exception.code, 0)
336        self.assertEqual(mock_stdout.getvalue(), "Interpreter: <none>\n")
337
338    def test_info_error(self):
339        # Test the info command fails when the archive does not exist.
340        target = self.tmpdir / 'dummy.pyz'
341        args = [str(target), '--info']
342        with self.assertRaises(SystemExit) as cm:
343            zipapp.main(args)
344        # Program should exit with a non-zero returm code.
345        self.assertTrue(cm.exception.code)
346
347
348if __name__ == "__main__":
349    unittest.main()
350