Ben Murdoch | 097c5b2 | 2016-05-18 11:27:45 +0100 | [diff] [blame^] | 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) |