gh-131178: Add tests for pstats command-line interface

The interactive profile browser started by ``python -m pstats`` had no
tests. Drive it through a subprocess and check each documented command,
the prompts, the report header, sorting, restrictions, callers and
callees, help, and the error paths.

Author: iamsharduld

Lib/test/test_pstats.py

@@ -8,8 +8,11 @@
 
 import os
 import pstats
+import re
+import subprocess
 import tempfile
 import cProfile
+from test.support import os_helper, script_helper
 
 class LazyImportTest(unittest.TestCase):
     @support.cpython_only
@@ -155,5 +158,222 @@ def test_SortKey_enum(self):
         self.assertEqual(SortKey.FILENAME, 'filename')
         self.assertNotEqual(SortKey.FILENAME, SortKey.CALLS)
 
+
+def _profiled_callee():
+    return 1
+
+
+def _profiled_caller():
+    for _ in range(3):
+        _profiled_callee()
+
+
+@support.requires_subprocess()
+class CommandLineTests(unittest.TestCase):
+    """Tests for the interactive profile browser (``python -m pstats``)."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.tmpdir = cls.enterClassContext(os_helper.temp_dir())
+        cls.profile = os.path.join(cls.tmpdir, 'test.prof')
+        cProfile.runctx('_profiled_caller()', globals(), {},
+                        filename=cls.profile)
+
+    def run_browser(self, *args, commands=()):
+        """Run ``python -m pstats`` with *args*, feeding *commands* on stdin.
+
+        Return what the browser wrote to stdout.
+        """
+        proc = script_helper.spawn_python('-m', 'pstats', *args,
+                                          stderr=subprocess.PIPE)
+        stdin = ''.join(f'{command}\n' for command in commands)
+        stdout, stderr = proc.communicate(stdin.encode())
+        stdout = stdout.decode('utf-8', 'backslashreplace')
+        stderr = stderr.decode('utf-8', 'backslashreplace')
+        self.assertEqual(proc.returncode, 0, stderr)
+        self.assertEqual(stderr, '')
+        return stdout.replace('\r\n', '\n')
+
+    def test_start_and_quit(self):
+        stdout = self.run_browser(commands=['quit'])
+        self.assertEqual(stdout,
+                         'Welcome to the profile statistics browser.\n'
+                         '% Goodbye.\n')
+
+    def test_eof(self):
+        # Closing stdin leaves the browser, like the "quit" command.
+        stdout = self.run_browser(commands=[])
+        self.assertEqual(stdout,
+                         'Welcome to the profile statistics browser.\n'
+                         '% \nGoodbye.\n')
+
+    def test_profile_argument(self):
+        stdout = self.run_browser(self.profile, commands=['stats', 'quit'])
+        # The prompt shows the name of the loaded file.
+        self.assertIn(f'\n{self.profile}% ', stdout)
+        # So does the header of the report.
+        self.assertIn(f'    {self.profile}\n', stdout)
+        self.assertRegex(stdout, r'\n +\d+ function calls in ')
+        self.assertRegex(stdout, r'\n +3 +.*\(_profiled_callee\)\n')
+        self.assertRegex(stdout, r'\n +1 +.*\(_profiled_caller\)\n')
+
+    def test_multiple_profile_arguments(self):
+        # The statistics of all the files given on the command line
+        # are added together.
+        stdout = self.run_browser(self.profile, self.profile,
+                                  commands=['stats', 'quit'])
+        self.assertEqual(stdout.count(f'    {self.profile}\n'), 2)
+        self.assertRegex(stdout, r'\n +6 +.*\(_profiled_callee\)\n')
+
+    def test_missing_profile_argument(self):
+        missing = os.path.join(self.tmpdir, 'missing.prof')
+        stdout = self.run_browser(missing, commands=['stats', 'quit'])
+        self.assertIn('No such file or directory\n', stdout)
+        self.assertIn('\n% No statistics object is loaded.\n', stdout)
+
+    def test_commands_without_statistics(self):
+        for command in ('stats', 'callers', 'callees', 'sort calls',
+                        'strip', 'reverse', f'add {self.profile}'):
+            with self.subTest(command=command):
+                stdout = self.run_browser(commands=[command, 'quit'])
+                self.assertIn('\n% No statistics object is loaded.\n',
+                              stdout)
+        stdout = self.run_browser(commands=['read', 'quit'])
+        self.assertIn(
+            '\n% No statistics object is current -- cannot reload.\n', stdout)
+
+    def test_read(self):
+        stdout = self.run_browser(commands=[f'read {self.profile}',
+                                            'stats', 'quit'])
+        self.assertIn(f'\n{self.profile}% ', stdout)
+        self.assertRegex(stdout, r'\n +3 +.*\(_profiled_callee\)\n')
+
+    def test_read_errors(self):
+        missing = os.path.join(self.tmpdir, 'missing.prof')
+        stdout = self.run_browser(commands=[f'read {missing}', 'quit'])
+        # The error is reported and no file is loaded.
+        self.assertEqual(stdout,
+                         'Welcome to the profile statistics browser.\n'
+                         '% No such file or directory\n'
+                         '% Goodbye.\n')
+
+        empty = os.path.join(self.tmpdir, 'empty.prof')
+        os_helper.create_empty_file(empty)
+        stdout = self.run_browser(commands=[f'read {empty}', 'quit'])
+        self.assertIn('\n% EOFError: ', stdout)
+        self.assertIn('\n% Goodbye.\n', stdout)
+
+    def test_add(self):
+        stdout = self.run_browser(self.profile,
+                                  commands=[f'add {self.profile}',
+                                            'stats', 'quit'])
+        self.assertEqual(stdout.count(f'    {self.profile}\n'), 2)
+        self.assertRegex(stdout, r'\n +6 +.*\(_profiled_callee\)\n')
+
+    def test_add_errors(self):
+        missing = os.path.join(self.tmpdir, 'missing.prof')
+        stdout = self.run_browser(self.profile,
+                                  commands=[f'add {missing}', 'quit'])
+        self.assertIn(f'% Failed to load statistics for {missing}: ', stdout)
+
+    def test_stats_restrictions(self):
+        stdout = self.run_browser(self.profile,
+                                  commands=['stats 2', 'stats 0.5',
+                                            'stats _profiled_callee', 'quit'])
+        self.assertRegex(stdout,
+                         r'List reduced from \d+ to 2 due to restriction <2>')
+        self.assertRegex(stdout, r'List reduced from \d+ to \d+ '
+                                 r'due to restriction <0\.5>')
+        self.assertRegex(stdout,
+                         r"List reduced from \d+ to 1 due to restriction "
+                         r"<'_profiled_callee'>")
+        by_name = stdout.split("<'_profiled_callee'>")[1]
+        self.assertIn('(_profiled_callee)', by_name)
+        self.assertNotIn('(_profiled_caller)', by_name)
+
+    def test_stats_invalid_fraction(self):
+        stdout = self.run_browser(self.profile,
+                                  commands=['stats 1.5', 'quit'])
+        self.assertIn('% Fraction argument must be in [0, 1]\n', stdout)
+        # The invalid argument is ignored and the full report is printed.
+        self.assertIn('(_profiled_caller)', stdout)
+        self.assertIn('(_profiled_callee)', stdout)
+
+    def test_sort(self):
+        # Unique prefixes of the sort keys are accepted.
+        stdout = self.run_browser(self.profile,
+                                  commands=['sort calls', 'stats',
+                                            'sort cumu', 'stats', 'quit'])
+        self.assertIn('Ordered by: call count', stdout)
+        self.assertIn('Ordered by: cumulative time', stdout)
+        by_calls = stdout.split('Ordered by: ')[1]
+        self.assertLess(by_calls.index('(_profiled_callee)'),
+                        by_calls.index('(_profiled_caller)'))
+
+    def test_sort_without_valid_keys(self):
+        for command in ('sort', 'sort bogus'):
+            with self.subTest(command=command):
+                stdout = self.run_browser(self.profile,
+                                          commands=[command, 'quit'])
+                self.assertIn(
+                    'Valid sort keys (unique prefixes are accepted):', stdout)
+                self.assertIn('\ncalls -- call count\n', stdout)
+                self.assertIn('\ncumulative -- cumulative time\n', stdout)
+
+    def test_reverse(self):
+        stdout = self.run_browser(self.profile,
+                                  commands=['sort calls', 'reverse',
+                                            'stats', 'quit'])
+        self.assertIn('Ordered by: call count', stdout)
+        self.assertLess(stdout.index('(_profiled_caller)'),
+                        stdout.index('(_profiled_callee)'))
+
+    def test_strip(self):
+        filename = _profiled_callee.__code__.co_filename
+        basename = os.path.basename(filename)
+        stdout = self.run_browser(self.profile,
+                                  commands=['stats _profiled_callee', 'strip',
+                                            'stats _profiled_callee', 'quit'])
+        before, after = stdout.split('due to restriction')[1:]
+        self.assertIn(f'{filename}:', before)
+        # Only the base name is left after "strip".
+        self.assertRegex(
+            after, r' ' + re.escape(basename) + r':\d+\(_profiled_callee\)')
+
+    def test_read_reloads_current_file(self):
+        # "read" without an argument reloads the current file, which
+        # undoes "strip".
+        filename = _profiled_callee.__code__.co_filename
+        stdout = self.run_browser(self.profile,
+                                  commands=['strip', 'read',
+                                            'stats _profiled_callee', 'quit'])
+        self.assertIn(f'{filename}:', stdout.split('due to restriction')[1])
+
+    def test_callers_and_callees(self):
+        stdout = self.run_browser(self.profile,
+                                  commands=['callers _profiled_callee',
+                                            'callees _profiled_caller',
+                                            'quit'])
+        self.assertIn('was called by...', stdout)
+        self.assertRegex(
+            stdout, r'\(_profiled_callee\) +<- +3 .*\(_profiled_caller\)\n')
+        self.assertIn('called...', stdout)
+        self.assertRegex(
+            stdout, r'\(_profiled_caller\) +-> +3 .*\(_profiled_callee\)\n')
+
+    def test_help(self):
+        stdout = self.run_browser(commands=['help', 'help stats', 'quit'])
+        lines = stdout.splitlines()
+        header = 'Documented commands (type help <topic>):'
+        self.assertIn(header, lines)
+        commands = lines[lines.index(header) + 2]
+        self.assertEqual(commands.split(),
+                         ['EOF', 'add', 'callees', 'callers', 'help', 'quit',
+                          'read', 'reverse', 'sort', 'stats', 'strip'])
+        self.assertIn('\n% Print statistics from the current stat object.\n',
+                      stdout)
+        self.assertIn('\nArguments may be:\n', stdout)
+
+
 if __name__ == "__main__":
     unittest.main()