build/android/pylib/base/test_collection.py

   1 # Copyright 2013 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 import threading
   6
   7 class TestCollection(object):
   8   """A threadsafe collection of tests.
   9
  10   Args:
  11     tests: List of tests to put in the collection.
  12   """
  13
  14   def __init__(self, tests=None):
  15     if not tests:
  16       tests = []
  17     self._lock = threading.Lock()
  18     self._tests = []
  19     self._tests_in_progress = 0
  20     # Used to signal that an item is available or all items have been handled.
  21     self._item_available_or_all_done = threading.Event()
  22     for t in tests:
  23       self.add(t)
  24
  25   def _pop(self):
  26     """Pop a test from the collection.
  27
  28     Waits until a test is available or all tests have been handled.
  29
  30     Returns:
  31       A test or None if all tests have been handled.
  32     """
  33     while True:
  34       # Wait for a test to be available or all tests to have been handled.
  35       self._item_available_or_all_done.wait()
  36       with self._lock:
  37         # Check which of the two conditions triggered the signal.
  38         if self._tests_in_progress == 0:
  39           return None
  40         try:
  41           return self._tests.pop(0)
  42         except IndexError:
  43           # Another thread beat us to the available test, wait again.
  44           self._item_available_or_all_done.clear()
  45
  46   def add(self, test):
  47     """Add a test to the collection.
  48
  49     Args:
  50       test: A test to add.
  51     """
  52     with self._lock:
  53       self._tests.append(test)
  54       self._item_available_or_all_done.set()
  55       self._tests_in_progress += 1
  56
  57   def test_completed(self):
  58     """Indicate that a test has been fully handled."""
  59     with self._lock:
  60       self._tests_in_progress -= 1
  61       if self._tests_in_progress == 0:
  62         # All tests have been handled, signal all waiting threads.
  63         self._item_available_or_all_done.set()
  64
  65   def __iter__(self):
  66     """Iterate through tests in the collection until all have been handled."""
  67     while True:
  68       r = self._pop()
  69       if r is None:
  70         break
  71       yield r
  72
  73   def __len__(self):
  74     """Return the number of tests currently in the collection."""
  75     return len(self._tests)
  76
  77   def test_names(self):
  78     """Return a list of the names of the tests currently in the collection."""
  79     with self._lock:
  80       return list(t.test for t in self._tests)