gh-131178: Add tests for `pickle` command-line interface (#131275)

Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>

2 files changed, 64 insertions, 5 deletions

diff --git a/Lib/pickle.py b/Lib/pickle.py
index efcdcbe..4fa3632 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -1907,7 +1907,7 @@ except ImportError:
     dump, dumps, load, loads = _dump, _dumps, _load, _loads
 
 
-if __name__ == "__main__":
+def _main(args=None):
     import argparse
     import pprint
     parser = argparse.ArgumentParser(
@@ -1915,7 +1915,7 @@ if __name__ == "__main__":
     parser.add_argument(
         'pickle_file',
         nargs='+', help='the pickle file')
-    args = parser.parse_args()
+    args = parser.parse_args(args)
     for fn in args.pickle_file:
         if fn == '-':
             obj = load(sys.stdin.buffer)
@@ -1923,3 +1923,7 @@ if __name__ == "__main__":
             with open(fn, 'rb') as f:
                 obj = load(f)
         pprint.pprint(obj)
+
+
+if __name__ == "__main__":
+    _main()
diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
index 385b257..296d4b8 100644
--- a/Lib/test/test_pickle.py
+++ b/Lib/test/test_pickle.py
@@ -1,18 +1,21 @@
 from _compat_pickle import (IMPORT_MAPPING, REVERSE_IMPORT_MAPPING,
                             NAME_MAPPING, REVERSE_NAME_MAPPING)
 import builtins
-import pickle
-import io
 import collections
+import contextlib
+import io
+import pickle
 import struct
 import sys
+import tempfile
 import warnings
 import weakref
+from textwrap import dedent
 
 import doctest
 import unittest
 from test import support
-from test.support import import_helper
+from test.support import import_helper, os_helper
 
 from test.pickletester import AbstractHookTests
 from test.pickletester import AbstractUnpickleTests
@@ -699,6 +702,58 @@ class CompatPickleTests(unittest.TestCase):
                                  ('multiprocessing.context', name))
 
 
+class CommandLineTest(unittest.TestCase):
+    def setUp(self):
+        self.filename = tempfile.mktemp()
+        self.addCleanup(os_helper.unlink, self.filename)
+
+    @staticmethod
+    def text_normalize(string):
+        """Dedent *string* and strip it from its surrounding whitespaces.
+
+        This method is used by the other utility functions so that any
+        string to write or to match against can be freely indented.
+        """
+        return dedent(string).strip()
+
+    def set_pickle_data(self, data):
+        with open(self.filename, 'wb') as f:
+            pickle.dump(data, f)
+
+    def invoke_pickle(self, *flags):
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            pickle._main(args=[*flags, self.filename])
+        return self.text_normalize(output.getvalue())
+
+    def test_invocation(self):
+        # test 'python -m pickle pickle_file'
+        data = {
+            'a': [1, 2.0, 3+4j],
+            'b': ('character string', b'byte string'),
+            'c': 'string'
+        }
+        expect = '''
+            {'a': [1, 2.0, (3+4j)],
+             'b': ('character string', b'byte string'),
+             'c': 'string'}
+        '''
+        self.set_pickle_data(data)
+
+        with self.subTest(data=data):
+            res = self.invoke_pickle()
+            expect = self.text_normalize(expect)
+            self.assertListEqual(res.splitlines(), expect.splitlines())
+
+    def test_unknown_flag(self):
+        stderr = io.StringIO()
+        with self.assertRaises(SystemExit):
+            # check that the parser help is shown
+            with contextlib.redirect_stderr(stderr):
+                _ = self.invoke_pickle('--unknown')
+        self.assertStartsWith(stderr.getvalue(), 'usage: ')
+
+
 def load_tests(loader, tests, pattern):
     tests.addTest(doctest.DocTestSuite(pickle))
     return tests