python/google/protobuf/internal/testing_refleaks.py

   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
  33 To use:
  34 - Use testing_refleak.BaseTestCase instead of unittest.TestCase
  35 - Configure and compile Python with --with-pydebug
  36
  37 If sys.gettotalrefcount() is not available (because Python was built without
  38 the Py_DEBUG option), then this module is a no-op and tests will run normally.
  39 """
  40
  41 import gc
  42 import sys
  43
  44 try:
  45   import copy_reg as copyreg  #PY26
  46 except ImportError:
  47   import copyreg
  48
  49 try:
  50   import unittest2 as unittest  #PY26
  51 except ImportError:
  52   import unittest
  53
  54
  55 class 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
  72 class 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
 114 if 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
 123 else:
 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