1# Copyright (c) 2012 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
6"""A module that contains a queue for running sharded tests."""
7
8import multiprocessing
9
10
11class ShardedTestsQueue(object):
12  """A queue for managing pending tests across different runners.
13
14  This class should only be used when sharding.
15
16  Attributes:
17    num_devices: an integer; the number of attached Android devices.
18    tests: a list of tests to be run.
19    tests_queue: if sharding, a JoinableQueue object that holds tests from
20        |tests|. Otherwise, a list holding tests.
21    results_queue: a Queue object to hold TestResults objects.
22  """
23  _STOP_SENTINEL = 'STOP'  # sentinel value for iter()
24
25  def __init__(self, num_devices, tests):
26    self.num_devices = num_devices
27    self.tests_queue = multiprocessing.Queue()
28    for test in tests:
29      self.tests_queue.put(test)
30    for _ in xrange(self.num_devices):
31      self.tests_queue.put(ShardedTestsQueue._STOP_SENTINEL)
32
33  def __iter__(self):
34    """Returns an iterator with the test cases."""
35    return iter(self.tests_queue.get, ShardedTestsQueue._STOP_SENTINEL)
36