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