• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import contextlib
2import os
3import os.path
4import sys
5import tempfile
6import test.support
7import unittest
8import unittest.mock
9
10import ensurepip
11import ensurepip._uninstall
12
13
14class TestPackages(unittest.TestCase):
15    def touch(self, directory, filename):
16        fullname = os.path.join(directory, filename)
17        open(fullname, "wb").close()
18
19    def test_version(self):
20        # Test version()
21        with tempfile.TemporaryDirectory() as tmpdir:
22            self.touch(tmpdir, "pip-1.2.3b1-py2.py3-none-any.whl")
23            self.touch(tmpdir, "setuptools-49.1.3-py3-none-any.whl")
24            with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None),
25                  unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir)):
26                self.assertEqual(ensurepip.version(), '1.2.3b1')
27
28    def test_get_packages_no_dir(self):
29        # Test _get_packages() without a wheel package directory
30        with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None),
31              unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', None)):
32            packages = ensurepip._get_packages()
33
34            # when bundled wheel packages are used, we get _PIP_VERSION
35            self.assertEqual(ensurepip._PIP_VERSION, ensurepip.version())
36
37        # use bundled wheel packages
38        self.assertIsNotNone(packages['pip'].wheel_name)
39        self.assertIsNotNone(packages['setuptools'].wheel_name)
40
41    def test_get_packages_with_dir(self):
42        # Test _get_packages() with a wheel package directory
43        setuptools_filename = "setuptools-49.1.3-py3-none-any.whl"
44        pip_filename = "pip-20.2.2-py2.py3-none-any.whl"
45
46        with tempfile.TemporaryDirectory() as tmpdir:
47            self.touch(tmpdir, setuptools_filename)
48            self.touch(tmpdir, pip_filename)
49            # not used, make sure that it's ignored
50            self.touch(tmpdir, "wheel-0.34.2-py2.py3-none-any.whl")
51
52            with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None),
53                  unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir)):
54                packages = ensurepip._get_packages()
55
56            self.assertEqual(packages['setuptools'].version, '49.1.3')
57            self.assertEqual(packages['setuptools'].wheel_path,
58                             os.path.join(tmpdir, setuptools_filename))
59            self.assertEqual(packages['pip'].version, '20.2.2')
60            self.assertEqual(packages['pip'].wheel_path,
61                             os.path.join(tmpdir, pip_filename))
62
63            # wheel package is ignored
64            self.assertEqual(sorted(packages), ['pip', 'setuptools'])
65
66
67class EnsurepipMixin:
68
69    def setUp(self):
70        run_pip_patch = unittest.mock.patch("ensurepip._run_pip")
71        self.run_pip = run_pip_patch.start()
72        self.run_pip.return_value = 0
73        self.addCleanup(run_pip_patch.stop)
74
75        # Avoid side effects on the actual os module
76        real_devnull = os.devnull
77        os_patch = unittest.mock.patch("ensurepip.os")
78        patched_os = os_patch.start()
79        # But expose os.listdir() used by _find_packages()
80        patched_os.listdir = os.listdir
81        self.addCleanup(os_patch.stop)
82        patched_os.devnull = real_devnull
83        patched_os.path = os.path
84        self.os_environ = patched_os.environ = os.environ.copy()
85
86
87class TestBootstrap(EnsurepipMixin, unittest.TestCase):
88
89    def test_basic_bootstrapping(self):
90        ensurepip.bootstrap()
91
92        self.run_pip.assert_called_once_with(
93            [
94                "install", "--no-cache-dir", "--no-index", "--find-links",
95                unittest.mock.ANY, "setuptools", "pip",
96            ],
97            unittest.mock.ANY,
98        )
99
100        additional_paths = self.run_pip.call_args[0][1]
101        self.assertEqual(len(additional_paths), 2)
102
103    def test_bootstrapping_with_root(self):
104        ensurepip.bootstrap(root="/foo/bar/")
105
106        self.run_pip.assert_called_once_with(
107            [
108                "install", "--no-cache-dir", "--no-index", "--find-links",
109                unittest.mock.ANY, "--root", "/foo/bar/",
110                "setuptools", "pip",
111            ],
112            unittest.mock.ANY,
113        )
114
115    def test_bootstrapping_with_user(self):
116        ensurepip.bootstrap(user=True)
117
118        self.run_pip.assert_called_once_with(
119            [
120                "install", "--no-cache-dir", "--no-index", "--find-links",
121                unittest.mock.ANY, "--user", "setuptools", "pip",
122            ],
123            unittest.mock.ANY,
124        )
125
126    def test_bootstrapping_with_upgrade(self):
127        ensurepip.bootstrap(upgrade=True)
128
129        self.run_pip.assert_called_once_with(
130            [
131                "install", "--no-cache-dir", "--no-index", "--find-links",
132                unittest.mock.ANY, "--upgrade", "setuptools", "pip",
133            ],
134            unittest.mock.ANY,
135        )
136
137    def test_bootstrapping_with_verbosity_1(self):
138        ensurepip.bootstrap(verbosity=1)
139
140        self.run_pip.assert_called_once_with(
141            [
142                "install", "--no-cache-dir", "--no-index", "--find-links",
143                unittest.mock.ANY, "-v", "setuptools", "pip",
144            ],
145            unittest.mock.ANY,
146        )
147
148    def test_bootstrapping_with_verbosity_2(self):
149        ensurepip.bootstrap(verbosity=2)
150
151        self.run_pip.assert_called_once_with(
152            [
153                "install", "--no-cache-dir", "--no-index", "--find-links",
154                unittest.mock.ANY, "-vv", "setuptools", "pip",
155            ],
156            unittest.mock.ANY,
157        )
158
159    def test_bootstrapping_with_verbosity_3(self):
160        ensurepip.bootstrap(verbosity=3)
161
162        self.run_pip.assert_called_once_with(
163            [
164                "install", "--no-cache-dir", "--no-index", "--find-links",
165                unittest.mock.ANY, "-vvv", "setuptools", "pip",
166            ],
167            unittest.mock.ANY,
168        )
169
170    def test_bootstrapping_with_regular_install(self):
171        ensurepip.bootstrap()
172        self.assertEqual(self.os_environ["ENSUREPIP_OPTIONS"], "install")
173
174    def test_bootstrapping_with_alt_install(self):
175        ensurepip.bootstrap(altinstall=True)
176        self.assertEqual(self.os_environ["ENSUREPIP_OPTIONS"], "altinstall")
177
178    def test_bootstrapping_with_default_pip(self):
179        ensurepip.bootstrap(default_pip=True)
180        self.assertNotIn("ENSUREPIP_OPTIONS", self.os_environ)
181
182    def test_altinstall_default_pip_conflict(self):
183        with self.assertRaises(ValueError):
184            ensurepip.bootstrap(altinstall=True, default_pip=True)
185        self.assertFalse(self.run_pip.called)
186
187    def test_pip_environment_variables_removed(self):
188        # ensurepip deliberately ignores all pip environment variables
189        # See http://bugs.python.org/issue19734 for details
190        self.os_environ["PIP_THIS_SHOULD_GO_AWAY"] = "test fodder"
191        ensurepip.bootstrap()
192        self.assertNotIn("PIP_THIS_SHOULD_GO_AWAY", self.os_environ)
193
194    def test_pip_config_file_disabled(self):
195        # ensurepip deliberately ignores the pip config file
196        # See http://bugs.python.org/issue20053 for details
197        ensurepip.bootstrap()
198        self.assertEqual(self.os_environ["PIP_CONFIG_FILE"], os.devnull)
199
200@contextlib.contextmanager
201def fake_pip(version=ensurepip.version()):
202    if version is None:
203        pip = None
204    else:
205        class FakePip():
206            __version__ = version
207        pip = FakePip()
208    sentinel = object()
209    orig_pip = sys.modules.get("pip", sentinel)
210    sys.modules["pip"] = pip
211    try:
212        yield pip
213    finally:
214        if orig_pip is sentinel:
215            del sys.modules["pip"]
216        else:
217            sys.modules["pip"] = orig_pip
218
219class TestUninstall(EnsurepipMixin, unittest.TestCase):
220
221    def test_uninstall_skipped_when_not_installed(self):
222        with fake_pip(None):
223            ensurepip._uninstall_helper()
224        self.assertFalse(self.run_pip.called)
225
226    def test_uninstall_skipped_with_warning_for_wrong_version(self):
227        with fake_pip("not a valid version"):
228            with test.support.captured_stderr() as stderr:
229                ensurepip._uninstall_helper()
230        warning = stderr.getvalue().strip()
231        self.assertIn("only uninstall a matching version", warning)
232        self.assertFalse(self.run_pip.called)
233
234
235    def test_uninstall(self):
236        with fake_pip():
237            ensurepip._uninstall_helper()
238
239        self.run_pip.assert_called_once_with(
240            [
241                "uninstall", "-y", "--disable-pip-version-check", "pip",
242                "setuptools",
243            ]
244        )
245
246    def test_uninstall_with_verbosity_1(self):
247        with fake_pip():
248            ensurepip._uninstall_helper(verbosity=1)
249
250        self.run_pip.assert_called_once_with(
251            [
252                "uninstall", "-y", "--disable-pip-version-check", "-v", "pip",
253                "setuptools",
254            ]
255        )
256
257    def test_uninstall_with_verbosity_2(self):
258        with fake_pip():
259            ensurepip._uninstall_helper(verbosity=2)
260
261        self.run_pip.assert_called_once_with(
262            [
263                "uninstall", "-y", "--disable-pip-version-check", "-vv", "pip",
264                "setuptools",
265            ]
266        )
267
268    def test_uninstall_with_verbosity_3(self):
269        with fake_pip():
270            ensurepip._uninstall_helper(verbosity=3)
271
272        self.run_pip.assert_called_once_with(
273            [
274                "uninstall", "-y", "--disable-pip-version-check", "-vvv",
275                "pip", "setuptools",
276            ]
277        )
278
279    def test_pip_environment_variables_removed(self):
280        # ensurepip deliberately ignores all pip environment variables
281        # See http://bugs.python.org/issue19734 for details
282        self.os_environ["PIP_THIS_SHOULD_GO_AWAY"] = "test fodder"
283        with fake_pip():
284            ensurepip._uninstall_helper()
285        self.assertNotIn("PIP_THIS_SHOULD_GO_AWAY", self.os_environ)
286
287    def test_pip_config_file_disabled(self):
288        # ensurepip deliberately ignores the pip config file
289        # See http://bugs.python.org/issue20053 for details
290        with fake_pip():
291            ensurepip._uninstall_helper()
292        self.assertEqual(self.os_environ["PIP_CONFIG_FILE"], os.devnull)
293
294
295# Basic testing of the main functions and their argument parsing
296
297EXPECTED_VERSION_OUTPUT = "pip " + ensurepip.version()
298
299class TestBootstrappingMainFunction(EnsurepipMixin, unittest.TestCase):
300
301    def test_bootstrap_version(self):
302        with test.support.captured_stdout() as stdout:
303            with self.assertRaises(SystemExit):
304                ensurepip._main(["--version"])
305        result = stdout.getvalue().strip()
306        self.assertEqual(result, EXPECTED_VERSION_OUTPUT)
307        self.assertFalse(self.run_pip.called)
308
309    def test_basic_bootstrapping(self):
310        exit_code = ensurepip._main([])
311
312        self.run_pip.assert_called_once_with(
313            [
314                "install", "--no-cache-dir", "--no-index", "--find-links",
315                unittest.mock.ANY, "setuptools", "pip",
316            ],
317            unittest.mock.ANY,
318        )
319
320        additional_paths = self.run_pip.call_args[0][1]
321        self.assertEqual(len(additional_paths), 2)
322        self.assertEqual(exit_code, 0)
323
324    def test_bootstrapping_error_code(self):
325        self.run_pip.return_value = 2
326        exit_code = ensurepip._main([])
327        self.assertEqual(exit_code, 2)
328
329
330class TestUninstallationMainFunction(EnsurepipMixin, unittest.TestCase):
331
332    def test_uninstall_version(self):
333        with test.support.captured_stdout() as stdout:
334            with self.assertRaises(SystemExit):
335                ensurepip._uninstall._main(["--version"])
336        result = stdout.getvalue().strip()
337        self.assertEqual(result, EXPECTED_VERSION_OUTPUT)
338        self.assertFalse(self.run_pip.called)
339
340    def test_basic_uninstall(self):
341        with fake_pip():
342            exit_code = ensurepip._uninstall._main([])
343
344        self.run_pip.assert_called_once_with(
345            [
346                "uninstall", "-y", "--disable-pip-version-check", "pip",
347                "setuptools",
348            ]
349        )
350
351        self.assertEqual(exit_code, 0)
352
353    def test_uninstall_error_code(self):
354        with fake_pip():
355            self.run_pip.return_value = 2
356            exit_code = ensurepip._uninstall._main([])
357        self.assertEqual(exit_code, 2)
358
359
360if __name__ == "__main__":
361    unittest.main()
362