summaryrefslogtreecommitdiff
path: root/unittest_utils.py
blob: e7090bcc6237e50d9d1f982c5db52f0e6f77bd0f (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#!/usr/bin/env python3

"""Helpers for unittests.  Note that when you import this we
automatically wrap unittest.main() with a call to bootstrap.initialize
so that we getLogger config, commandline args, logging control,
etc... this works fine but it's a little hacky so caveat emptor.
"""

import functools
import inspect
import io
import logging
import pickle
import random
import statistics
import sys
import time
import tempfile
from typing import Callable, Iterable
import unittest

import bootstrap
import config


logger = logging.getLogger(__name__)
cfg = config.add_commandline_args(
    f'Logging ({__file__})',
    'Args related to function decorators')
cfg.add_argument(
    '--unittests_ignore_perf',
    action='store_true',
    default=False,
    help='Ignore unittest perf regression in @check_method_for_perf_regressions',
)
cfg.add_argument(
    '--unittests_num_perf_samples',
    type=int,
    default=20,
    help='The count of perf timing samples we need to see before blocking slow runs on perf grounds'
)
cfg.add_argument(
    '--unittests_drop_perf_traces',
    type=str,
    nargs=1,
    default=None,
    help='The identifier (i.e. file!test_fixture) for which we should drop all perf data'
)


# >>> This is the hacky business, FYI. <<<
unittest.main = bootstrap.initialize(unittest.main)


_db = '/home/scott/.python_unittest_performance_db'


def check_method_for_perf_regressions(func: Callable) -> Callable:
    """This is meant to be used on a method in a class that subclasses
    unittest.TestCase.  When thus decorated it will time the execution
    of the code in the method, compare it with a database of
    historical perfmance, and fail the test with a perf-related
    message if it has become too slow.
    """

    def load_known_test_performance_characteristics():
        with open(_db, 'rb') as f:
            return pickle.load(f)

    def save_known_test_performance_characteristics(perfdb):
        with open(_db, 'wb') as f:
            pickle.dump(perfdb, f, pickle.HIGHEST_PROTOCOL)

    @functools.wraps(func)
    def wrapper_perf_monitor(*args, **kwargs):
        try:
            perfdb = load_known_test_performance_characteristics()
        except Exception as e:
            logger.exception(e)
            logger.warning(f'Unable to load perfdb from {_db}')
            perfdb = {}

        # This is a unique identifier for a test: filepath!function
        logger.debug(f'Watching {func.__name__}\'s performance...')
        func_id = f'{func.__globals__["__file__"]}!{func.__name__}'
        logger.debug(f'Canonical function identifier = {func_id}')

        # cmdline arg to forget perf traces for function
        drop_id = config.config['unittests_drop_perf_traces']
        if drop_id is not None:
            if drop_id in perfdb:
                perfdb[drop_id] = []

        # Run the wrapped test paying attention to latency.
        start_time = time.perf_counter()
        value = func(*args, **kwargs)
        end_time = time.perf_counter()
        run_time = end_time - start_time
        logger.debug(f'{func.__name__} executed in {run_time:f}s.')

        # Check the db; see if it was unexpectedly slow.
        hist = perfdb.get(func_id, [])
        if len(hist) < config.config['unittests_num_perf_samples']:
            hist.append(run_time)
            logger.debug(
                f'Still establishing a perf baseline for {func.__name__}'
            )
        else:
            stdev = statistics.stdev(hist)
            limit = hist[-1] + stdev * 3
            logger.debug(
                f'Max acceptable performace for {func.__name__} is {limit:f}s'
            )
            if (
                run_time > limit and
                not config.config['unittests_ignore_perf']
            ):
                msg = f'''{func_id} performance has regressed unacceptably.
{hist[-1]:f}s is the slowest record in {len(hist)} db perf samples.
It just ran in {run_time:f}s which is >3 stdevs slower than the slowest sample.
Here is the current, full db perf timing distribution:

{hist}'''
                slf = args[0]
                logger.error(msg)
                slf.fail(msg)
            else:
                hist.append(run_time)

        n = min(config.config['unittests_num_perf_samples'], len(hist))
        hist = random.sample(hist, n)
        hist.sort()
        perfdb[func_id] = hist
        save_known_test_performance_characteristics(perfdb)
        return value
    return wrapper_perf_monitor


def check_all_methods_for_perf_regressions(prefix='test_'):
    def decorate_the_testcase(cls):
        if issubclass(cls, unittest.TestCase):
            for name, m in inspect.getmembers(cls, inspect.isfunction):
                if name.startswith(prefix):
                    setattr(cls, name, check_method_for_perf_regressions(m))
                    logger.debug(f'Wrapping {cls.__name__}:{name}.')
        return cls
    return decorate_the_testcase


def breakpoint():
    import pdb
    pdb.set_trace()