testing_support/thread_watcher.py

# Copyright (c) 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.


import logging
import sys
import threading
import traceback
import unittest


class ThreadWatcherMixIn(object):
  def setUp(self):
    self._pre_test_threads = [t.ident for t in threading.enumerate()]

  def tearDown(self):
    # If test failed before this check, don't raise another exception
    # overwriting the original one, and making it harder to debug.
    # Besides, that exception might be the reason cleanup didn't happen.
    if self._test_definitely_failed():
      return

    post_test_threads = threading.enumerate()
    new_threads = [t for t in post_test_threads
                   if t.ident not in self._pre_test_threads]
    if new_threads:
      details = []
      for th in new_threads:
        details.append('\nThread %s stacktrace:\n' % th)
        try:
          details.extend(traceback.format_stack(sys._current_frames()[t.ident]))
        except Exception:
          if t.ident in sys._current_frames():
            raise
          # This can happen due to threads starting or stopping concurrently.
          details.append('  Thread stopped while acquiring stacktrace.\n')
      details = ''.join(details)

      self.fail('Found %d running thread(s) after the test.\n%s' % (
                len(new_threads), details))

  def _test_definitely_failed(self):
    """Returns True only if test has definitely failed already.

    This is necessary to determine whether we should fail the test when finding
    stray threads. If the test has failed in some other way, we should avoid
    throwing an exception because it will mask the original error.
    """
    # When run using unittest runner,
    # self._resultForDoCleanups is set to internal unittest object.
    if hasattr(self, '_resultForDoCleanups') and self._resultForDoCleanups:
      return not self._resultForDoCleanups.wasSuccessful()
    # expect_tests runner does so differently (see
    # https://codereview.chromium.org/2121343004).
    if hasattr(self, '_test_failed_with_exception'):
      return self._test_failed_with_exception
    # Default case - be conservative.
    return False


class TestCase(unittest.TestCase, ThreadWatcherMixIn):
  """Base unittest class that fails on leaked threads after the test."""
  def setUp(self):
    super(TestCase, self).setUp()
    ThreadWatcherMixIn.setUp(self)

  def tearDown(self):
    ThreadWatcherMixIn.tearDown(self)
    super(TestCase, self).tearDown()