1  import unittest
       2  
       3  from test import support
       4  from io import StringIO
       5  from pstats import SortKey
       6  from enum import StrEnum, _test_simple_enum
       7  
       8  import pstats
       9  import cProfile
      10  
      11  class ESC[4;38;5;81mAddCallersTestCase(ESC[4;38;5;149munittestESC[4;38;5;149m.ESC[4;38;5;149mTestCase):
      12      """Tests for pstats.add_callers helper."""
      13  
      14      def test_combine_results(self):
      15          # pstats.add_callers should combine the call results of both target
      16          # and source by adding the call time. See issue1269.
      17          # new format: used by the cProfile module
      18          target = {"a": (1, 2, 3, 4)}
      19          source = {"a": (1, 2, 3, 4), "b": (5, 6, 7, 8)}
      20          new_callers = pstats.add_callers(target, source)
      21          self.assertEqual(new_callers, {'a': (2, 4, 6, 8), 'b': (5, 6, 7, 8)})
      22          # old format: used by the profile module
      23          target = {"a": 1}
      24          source = {"a": 1, "b": 5}
      25          new_callers = pstats.add_callers(target, source)
      26          self.assertEqual(new_callers, {'a': 2, 'b': 5})
      27  
      28  
      29  class ESC[4;38;5;81mStatsTestCase(ESC[4;38;5;149munittestESC[4;38;5;149m.ESC[4;38;5;149mTestCase):
      30      def setUp(self):
      31          stats_file = support.findfile('pstats.pck')
      32          self.stats = pstats.Stats(stats_file)
      33  
      34      def test_add(self):
      35          stream = StringIO()
      36          stats = pstats.Stats(stream=stream)
      37          stats.add(self.stats, self.stats)
      38  
      39      def test_sort_stats_int(self):
      40          valid_args = {-1: 'stdname',
      41                        0: 'calls',
      42                        1: 'time',
      43                        2: 'cumulative'}
      44          for arg_int, arg_str in valid_args.items():
      45              self.stats.sort_stats(arg_int)
      46              self.assertEqual(self.stats.sort_type,
      47                               self.stats.sort_arg_dict_default[arg_str][-1])
      48  
      49      def test_sort_stats_string(self):
      50          for sort_name in ['calls', 'ncalls', 'cumtime', 'cumulative',
      51                      'filename', 'line', 'module', 'name', 'nfl', 'pcalls',
      52                      'stdname', 'time', 'tottime']:
      53              self.stats.sort_stats(sort_name)
      54              self.assertEqual(self.stats.sort_type,
      55                               self.stats.sort_arg_dict_default[sort_name][-1])
      56  
      57      def test_sort_stats_partial(self):
      58          sortkey = 'filename'
      59          for sort_name in ['f', 'fi', 'fil', 'file', 'filen', 'filena',
      60                             'filenam', 'filename']:
      61              self.stats.sort_stats(sort_name)
      62              self.assertEqual(self.stats.sort_type,
      63                               self.stats.sort_arg_dict_default[sortkey][-1])
      64  
      65      def test_sort_stats_enum(self):
      66          for member in SortKey:
      67              self.stats.sort_stats(member)
      68              self.assertEqual(
      69                      self.stats.sort_type,
      70                      self.stats.sort_arg_dict_default[member.value][-1])
      71          class ESC[4;38;5;81mCheckedSortKey(ESC[4;38;5;149mStrEnum):
      72              CALLS = 'calls', 'ncalls'
      73              CUMULATIVE = 'cumulative', 'cumtime'
      74              FILENAME = 'filename', 'module'
      75              LINE = 'line'
      76              NAME = 'name'
      77              NFL = 'nfl'
      78              PCALLS = 'pcalls'
      79              STDNAME = 'stdname'
      80              TIME = 'time', 'tottime'
      81              def __new__(cls, *values):
      82                  value = values[0]
      83                  obj = str.__new__(cls, value)
      84                  obj._value_ = value
      85                  for other_value in values[1:]:
      86                      cls._value2member_map_[other_value] = obj
      87                  obj._all_values = values
      88                  return obj
      89          _test_simple_enum(CheckedSortKey, SortKey)
      90  
      91      def test_sort_starts_mix(self):
      92          self.assertRaises(TypeError, self.stats.sort_stats,
      93                            'calls',
      94                            SortKey.TIME)
      95          self.assertRaises(TypeError, self.stats.sort_stats,
      96                            SortKey.TIME,
      97                            'calls')
      98  
      99      def test_get_stats_profile(self):
     100          def pass1(): pass
     101          def pass2(): pass
     102          def pass3(): pass
     103  
     104          pr = cProfile.Profile()
     105          pr.enable()
     106          pass1()
     107          pass2()
     108          pass3()
     109          pr.create_stats()
     110          ps = pstats.Stats(pr)
     111  
     112          stats_profile = ps.get_stats_profile()
     113          funcs_called = set(stats_profile.func_profiles.keys())
     114          self.assertIn('pass1', funcs_called)
     115          self.assertIn('pass2', funcs_called)
     116          self.assertIn('pass3', funcs_called)
     117  
     118      def test_SortKey_enum(self):
     119          self.assertEqual(SortKey.FILENAME, 'filename')
     120          self.assertNotEqual(SortKey.FILENAME, SortKey.CALLS)
     121  
     122  if __name__ == "__main__":
     123      unittest.main()