1import unittest
2
3from test import support
4from io import StringIO
5from pstats import SortKey
6from enum import StrEnum, _test_simple_enum
7
8import pstats
9import cProfile
10
11class AddCallersTestCase(unittest.TestCase):
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
29class StatsTestCase(unittest.TestCase):
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 CheckedSortKey(StrEnum):
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
122if __name__ == "__main__":
123    unittest.main()
124