1# Protocol Buffers - Google's data interchange format
2# Copyright 2008 Google Inc.  All rights reserved.
3# https://developers.google.com/protocol-buffers/
4#
5# Redistribution and use in source and binary forms, with or without
6# modification, are permitted provided that the following conditions are
7# met:
8#
9#     * Redistributions of source code must retain the above copyright
10# notice, this list of conditions and the following disclaimer.
11#     * Redistributions in binary form must reproduce the above
12# copyright notice, this list of conditions and the following disclaimer
13# in the documentation and/or other materials provided with the
14# distribution.
15#     * Neither the name of Google Inc. nor the names of its
16# contributors may be used to endorse or promote products derived from
17# this software without specific prior written permission.
18#
19# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
22# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
23# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
26# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
27# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
28# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30
31"""A subclass of unittest.TestCase which checks for reference leaks.
32
33To use:
34- Use testing_refleak.BaseTestCase instead of unittest.TestCase
35- Configure and compile Python with --with-pydebug
36
37If sys.gettotalrefcount() is not available (because Python was built without
38the Py_DEBUG option), then this module is a no-op and tests will run normally.
39"""
40
41import gc
42import sys
43
44try:
45  import copy_reg as copyreg  #PY26
46except ImportError:
47  import copyreg
48
49try:
50  import unittest2 as unittest  #PY26
51except ImportError:
52  import unittest
53
54
55class LocalTestResult(unittest.TestResult):
56  """A TestResult which forwards events to a parent object, except for Skips."""
57
58  def __init__(self, parent_result):
59    unittest.TestResult.__init__(self)
60    self.parent_result = parent_result
61
62  def addError(self, test, error):
63    self.parent_result.addError(test, error)
64
65  def addFailure(self, test, error):
66    self.parent_result.addFailure(test, error)
67
68  def addSkip(self, test, reason):
69    pass
70
71
72class ReferenceLeakCheckerMixin(object):
73  """A mixin class for TestCase, which checks reference counts."""
74
75  NB_RUNS = 3
76
77  def run(self, result=None):
78    # python_message.py registers all Message classes to some pickle global
79    # registry, which makes the classes immortal.
80    # We save a copy of this registry, and reset it before we could references.
81    self._saved_pickle_registry = copyreg.dispatch_table.copy()
82
83    # Run the test twice, to warm up the instance attributes.
84    super(ReferenceLeakCheckerMixin, self).run(result=result)
85    super(ReferenceLeakCheckerMixin, self).run(result=result)
86
87    oldrefcount = 0
88    local_result = LocalTestResult(result)
89
90    refcount_deltas = []
91    for _ in range(self.NB_RUNS):
92      oldrefcount = self._getRefcounts()
93      super(ReferenceLeakCheckerMixin, self).run(result=local_result)
94      newrefcount = self._getRefcounts()
95      refcount_deltas.append(newrefcount - oldrefcount)
96    print(refcount_deltas, self)
97
98    try:
99      self.assertEqual(refcount_deltas, [0] * self.NB_RUNS)
100    except Exception:  # pylint: disable=broad-except
101      result.addError(self, sys.exc_info())
102
103  def _getRefcounts(self):
104    copyreg.dispatch_table.clear()
105    copyreg.dispatch_table.update(self._saved_pickle_registry)
106    # It is sometimes necessary to gc.collect() multiple times, to ensure
107    # that all objects can be collected.
108    gc.collect()
109    gc.collect()
110    gc.collect()
111    return sys.gettotalrefcount()
112
113
114if hasattr(sys, 'gettotalrefcount'):
115
116  def TestCase(test_class):
117    new_bases = (ReferenceLeakCheckerMixin,) + test_class.__bases__
118    new_class = type(test_class)(
119        test_class.__name__, new_bases, dict(test_class.__dict__))
120    return new_class
121  SkipReferenceLeakChecker = unittest.skip
122
123else:
124  # When PyDEBUG is not enabled, run the tests normally.
125
126  def TestCase(test_class):
127    return test_class
128
129  def SkipReferenceLeakChecker(reason):
130    del reason  # Don't skip, so don't need a reason.
131    def Same(func):
132      return func
133    return Same
134