Issue #19728: fix ensurepip name clash with submodule

Also added refactoring and added basic tests for the argument
parsing in both ensurepip._main and ensurepip._uninstall._main.

4 files changed, 127 insertions, 91 deletions

diff --git a/Lib/ensurepip/__init__.py b/Lib/ensurepip/__init__.py
index 090c410..0f43bc4 100644
--- a/Lib/ensurepip/__init__.py
+++ b/Lib/ensurepip/__init__.py
@@ -104,7 +104,7 @@ def bootstrap(*, root=None, upgrade=False, user=False,
 
         _run_pip(args + [p[0] for p in _PROJECTS], additional_paths)
 
-def _uninstall(*, verbosity=0):
+def _uninstall_helper(*, verbosity=0):
     """Helper to support a clean default uninstall process on Windows
 
     Note that calling this function may alter os.environ.
@@ -129,3 +129,64 @@ def _uninstall(*, verbosity=0):
         args += ["-" + "v" * verbosity]
 
     _run_pip(args + [p[0] for p in reversed(_PROJECTS)])
+
+
+def _main(argv=None):
+    import argparse
+    parser = argparse.ArgumentParser(prog="python -m ensurepip")
+    parser.add_argument(
+        "--version",
+        action="version",
+        version="pip {}".format(version()),
+        help="Show the version of pip that is bundled with this Python.",
+    )
+    parser.add_argument(
+        "-v", "--verbose",
+        action="count",
+        default=0,
+        dest="verbosity",
+        help=("Give more output. Option is additive, and can be used up to 3 "
+              "times."),
+    )
+    parser.add_argument(
+        "-U", "--upgrade",
+        action="store_true",
+        default=False,
+        help="Upgrade pip and dependencies, even if already installed.",
+    )
+    parser.add_argument(
+        "--user",
+        action="store_true",
+        default=False,
+        help="Install using the user scheme.",
+    )
+    parser.add_argument(
+        "--root",
+        default=None,
+        help="Install everything relative to this alternate root directory.",
+    )
+    parser.add_argument(
+        "--altinstall",
+        action="store_true",
+        default=False,
+        help=("Make an alternate install, installing only the X.Y versioned"
+              "scripts (Default: pipX, pipX.Y, easy_install-X.Y)"),
+    )
+    parser.add_argument(
+        "--default-pip",
+        action="store_true",
+        default=False,
+        help=("Make a default pip install, installing the unqualified pip "
+              "and easy_install in addition to the versioned scripts"),
+    )
+
+    args = parser.parse_args(argv)
+
+    bootstrap(
+        root=args.root,
+        upgrade=args.upgrade,
+        user=args.user,
+        verbosity=args.verbosity,
+        altinstall=args.altinstall,
+        default_pip=args.default_pip,
+    )
diff --git a/Lib/ensurepip/__main__.py b/Lib/ensurepip/__main__.py
index 53e8459..77527d7 100644
--- a/Lib/ensurepip/__main__.py
+++ b/Lib/ensurepip/__main__.py
@@ -1,66 +1,4 @@
-import argparse
 import ensurepip
 
-
-def main():
-    parser = argparse.ArgumentParser(prog="python -m ensurepip")
-    parser.add_argument(
-        "--version",
-        action="version",
-        version="pip {}".format(ensurepip.version()),
-        help="Show the version of pip that is bundled with this Python.",
-    )
-    parser.add_argument(
-        "-v", "--verbose",
-        action="count",
-        default=0,
-        dest="verbosity",
-        help=("Give more output. Option is additive, and can be used up to 3 "
-              "times."),
-    )
-    parser.add_argument(
-        "-U", "--upgrade",
-        action="store_true",
-        default=False,
-        help="Upgrade pip and dependencies, even if already installed.",
-    )
-    parser.add_argument(
-        "--user",
-        action="store_true",
-        default=False,
-        help="Install using the user scheme.",
-    )
-    parser.add_argument(
-        "--root",
-        default=None,
-        help="Install everything relative to this alternate root directory.",
-    )
-    parser.add_argument(
-        "--altinstall",
-        action="store_true",
-        default=False,
-        help=("Make an alternate install, installing only the X.Y versioned"
-              "scripts (Default: pipX, pipX.Y, easy_install-X.Y)"),
-    )
-    parser.add_argument(
-        "--default-pip",
-        action="store_true",
-        default=False,
-        help=("Make a default pip install, installing the unqualified pip "
-              "and easy_install in addition to the versioned scripts"),
-    )
-
-    args = parser.parse_args()
-
-    ensurepip.bootstrap(
-        root=args.root,
-        upgrade=args.upgrade,
-        user=args.user,
-        verbosity=args.verbosity,
-        altinstall=args.altinstall,
-        default_pip=args.default_pip,
-    )
-
-
 if __name__ == "__main__":
-    main()
+    ensurepip._main()
diff --git a/Lib/ensurepip/_uninstall.py b/Lib/ensurepip/_uninstall.py
index 38c486b..750365e 100644
--- a/Lib/ensurepip/_uninstall.py
+++ b/Lib/ensurepip/_uninstall.py
@@ -4,7 +4,7 @@ import argparse
 import ensurepip
 
 
-def main():
+def _main(argv=None):
     parser = argparse.ArgumentParser(prog="python -m ensurepip._uninstall")
     parser.add_argument(
         "--version",
@@ -21,10 +21,10 @@ def main():
               "times."),
     )
 
-    args = parser.parse_args()
+    args = parser.parse_args(argv)
 
-    ensurepip._uninstall(verbosity=args.verbosity)
+    ensurepip._uninstall_helper(verbosity=args.verbosity)
 
 
 if __name__ == "__main__":
-    main()
+    _main()
diff --git a/Lib/test/test_ensurepip.py b/Lib/test/test_ensurepip.py
index 2080ef1..ad2311d 100644
--- a/Lib/test/test_ensurepip.py
+++ b/Lib/test/test_ensurepip.py
@@ -1,20 +1,20 @@
 import unittest
 import unittest.mock
-import ensurepip
 import test.support
 import os
 import os.path
 import contextlib
 import sys
 
+import ensurepip
+import ensurepip._uninstall
 
 class TestEnsurePipVersion(unittest.TestCase):
 
     def test_returns_version(self):
         self.assertEqual(ensurepip._PIP_VERSION, ensurepip.version())
 
-
-class TestBootstrap(unittest.TestCase):
+class EnsurepipMixin:
 
     def setUp(self):
         run_pip_patch = unittest.mock.patch("ensurepip._run_pip")
@@ -28,6 +28,8 @@ class TestBootstrap(unittest.TestCase):
         patched_os.path = os.path
         self.os_environ = patched_os.environ = os.environ.copy()
 
+class TestBootstrap(EnsurepipMixin, unittest.TestCase):
+
     def test_basic_bootstrapping(self):
         ensurepip.bootstrap()
 
@@ -153,35 +155,23 @@ def fake_pip(version=ensurepip._PIP_VERSION):
         else:
             sys.modules["pip"] = orig_pip
 
-class TestUninstall(unittest.TestCase):
-
-    def setUp(self):
-        run_pip_patch = unittest.mock.patch("ensurepip._run_pip")
-        self.run_pip = run_pip_patch.start()
-        self.addCleanup(run_pip_patch.stop)
-
-        # Avoid side effects on the actual os module
-        os_patch = unittest.mock.patch("ensurepip.os")
-        patched_os = os_patch.start()
-        self.addCleanup(os_patch.stop)
-        patched_os.path = os.path
-        self.os_environ = patched_os.environ = os.environ.copy()
+class TestUninstall(EnsurepipMixin, unittest.TestCase):
 
     def test_uninstall_skipped_when_not_installed(self):
         with fake_pip(None):
-            ensurepip._uninstall()
+            ensurepip._uninstall_helper()
         self.run_pip.assert_not_called()
 
     def test_uninstall_fails_with_wrong_version(self):
         with fake_pip("not a valid version"):
             with self.assertRaises(RuntimeError):
-                ensurepip._uninstall()
+                ensurepip._uninstall_helper()
         self.run_pip.assert_not_called()
 
 
     def test_uninstall(self):
         with fake_pip():
-            ensurepip._uninstall()
+            ensurepip._uninstall_helper()
 
         self.run_pip.assert_called_once_with(
             ["uninstall", "-y", "pip", "setuptools"]
@@ -189,7 +179,7 @@ class TestUninstall(unittest.TestCase):
 
     def test_uninstall_with_verbosity_1(self):
         with fake_pip():
-            ensurepip._uninstall(verbosity=1)
+            ensurepip._uninstall_helper(verbosity=1)
 
         self.run_pip.assert_called_once_with(
             ["uninstall", "-y", "-v", "pip", "setuptools"]
@@ -197,7 +187,7 @@ class TestUninstall(unittest.TestCase):
 
     def test_uninstall_with_verbosity_2(self):
         with fake_pip():
-            ensurepip._uninstall(verbosity=2)
+            ensurepip._uninstall_helper(verbosity=2)
 
         self.run_pip.assert_called_once_with(
             ["uninstall", "-y", "-vv", "pip", "setuptools"]
@@ -205,7 +195,7 @@ class TestUninstall(unittest.TestCase):
 
     def test_uninstall_with_verbosity_3(self):
         with fake_pip():
-            ensurepip._uninstall(verbosity=3)
+            ensurepip._uninstall_helper(verbosity=3)
 
         self.run_pip.assert_called_once_with(
             ["uninstall", "-y", "-vvv", "pip", "setuptools"]
@@ -216,10 +206,57 @@ class TestUninstall(unittest.TestCase):
         # See http://bugs.python.org/issue19734 for details
         self.os_environ["PIP_THIS_SHOULD_GO_AWAY"] = "test fodder"
         with fake_pip():
-            ensurepip._uninstall()
+            ensurepip._uninstall_helper()
         self.assertNotIn("PIP_THIS_SHOULD_GO_AWAY", self.os_environ)
 
 
+# Basic testing of the main functions and their argument parsing
+
+EXPECTED_VERSION_OUTPUT = "pip " + ensurepip._PIP_VERSION
+
+class TestBootstrappingMainFunction(EnsurepipMixin, unittest.TestCase):
+
+    def test_bootstrap_version(self):
+        with test.support.captured_stdout() as stdout:
+            with self.assertRaises(SystemExit):
+                ensurepip._main(["--version"])
+        result = stdout.getvalue().strip()
+        self.assertEqual(result, EXPECTED_VERSION_OUTPUT)
+        self.run_pip.assert_not_called()
+
+    def test_basic_bootstrapping(self):
+        ensurepip._main([])
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-index", "--find-links",
+                unittest.mock.ANY, "--pre", "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+        additional_paths = self.run_pip.call_args[0][1]
+        self.assertEqual(len(additional_paths), 2)
+
+class TestUninstallationMainFunction(EnsurepipMixin, unittest.TestCase):
+
+    def test_uninstall_version(self):
+        with test.support.captured_stdout() as stdout:
+            with self.assertRaises(SystemExit):
+                ensurepip._uninstall._main(["--version"])
+        result = stdout.getvalue().strip()
+        self.assertEqual(result, EXPECTED_VERSION_OUTPUT)
+        self.run_pip.assert_not_called()
+
+    def test_basic_uninstall(self):
+        with fake_pip():
+            ensurepip._uninstall._main([])
+
+        self.run_pip.assert_called_once_with(
+            ["uninstall", "-y", "pip", "setuptools"]
+        )
+
+
 
 if __name__ == "__main__":
     test.support.run_unittest(__name__)