michael@0: # Copyright (c) 2012 The Chromium Authors. All rights reserved.
michael@0: # Use of this source code is governed by a BSD-style license that can be
michael@0: # found in the LICENSE file.
michael@0: 
michael@0: 
michael@0: """A module that contains a queue for running sharded tests."""
michael@0: 
michael@0: import multiprocessing
michael@0: 
michael@0: 
michael@0: class ShardedTestsQueue(object):
michael@0:   """A queue for managing pending tests across different runners.
michael@0: 
michael@0:   This class should only be used when sharding.
michael@0: 
michael@0:   Attributes:
michael@0:     num_devices: an integer; the number of attached Android devices.
michael@0:     tests: a list of tests to be run.
michael@0:     tests_queue: if sharding, a JoinableQueue object that holds tests from
michael@0:         |tests|. Otherwise, a list holding tests.
michael@0:     results_queue: a Queue object to hold TestResults objects.
michael@0:   """
michael@0:   _STOP_SENTINEL = 'STOP'  # sentinel value for iter()
michael@0: 
michael@0:   def __init__(self, num_devices, tests):
michael@0:     self.num_devices = num_devices
michael@0:     self.tests_queue = multiprocessing.Queue()
michael@0:     for test in tests:
michael@0:       self.tests_queue.put(test)
michael@0:     for _ in xrange(self.num_devices):
michael@0:       self.tests_queue.put(ShardedTestsQueue._STOP_SENTINEL)
michael@0: 
michael@0:   def __iter__(self):
michael@0:     """Returns an iterator with the test cases."""
michael@0:     return iter(self.tests_queue.get, ShardedTestsQueue._STOP_SENTINEL)