1# Copyright 2015 The Chromium OS 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
6class DroneTaskQueueException(Exception):
7    """Generic task queue exception."""
8    pass
9
10
11class DroneTaskQueue(object):
12    """A manager to run queued tasks in drones and gather results from them."""
13
14    def __init__(self):
15        self.results = dict()
16
17
18    def get_results(self):
19        """Get a results dictionary keyed on drones.
20
21        @return: A dictionary of return values from drones.
22        """
23        results_copy = self.results.copy()
24        self.results.clear()
25        return results_copy
26
27
28    def execute(self, drones, wait=True):
29        """Invoke methods via SSH to a drone.
30
31        @param drones: A list of drones with calls to execute.
32        @param wait: If True, this method will only return when all the drones
33            have returned the result of their respective invocations of
34            drone_utility. The `results` map will be cleared.
35            If False, the caller must clear the map before the next invocation
36            of `execute`, by calling `get_results`.
37
38        @return: A dictionary keyed on the drones, containing a list of return
39            values from the execution of drone_utility.
40
41        @raises DroneTaskQueueException: If the results map isn't empty at the
42            time of invocation.
43        """
44        if self.results:
45            raise DroneTaskQueueException(
46                    'Cannot clobber results map: %s, it should be cleared '
47                    'through get_results.' % self.results)
48        for drone in drones:
49            if not drone.get_calls():
50                continue
51            drone_results = drone.execute_queued_calls()
52            if drone in self.results:
53                raise DroneTaskQueueException(
54                        'Task queue has recorded results for drone %s: %s' %
55                        (drone, self.results))
56            self.results[drone] = drone_results
57        return self.get_results() if wait else None
58