1#!/usr/bin/python
2#
3# Copyright 2015 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""Tests for the drone task queue."""
8
9import cPickle
10import logging
11import Queue
12import unittest
13
14import common
15from autotest_lib.client.common_lib import utils
16from autotest_lib.client.common_lib.test_utils import mock
17from autotest_lib.scheduler import drone_task_queue
18from autotest_lib.scheduler import drones
19from autotest_lib.server.hosts import ssh_host
20
21
22class DroneTaskQueueTest(unittest.TestCase):
23    """Drone task queue tests."""
24
25    def create_remote_drone(self, hostname):
26        """Create and initialize a Remote Drone.
27
28        @param hostname: The name of the host for the remote drone.
29
30        @return: A remote drone instance.
31        """
32        drones.drone_utility.create_host.expect_call(hostname).and_return(
33                self._mock_host)
34        self._mock_host.is_up.expect_call().and_return(True)
35        return drones._RemoteDrone(hostname, timestamp_remote_calls=False)
36
37
38    def setUp(self):
39        self.god = mock.mock_god()
40        self._mock_host = self.god.create_mock_class(ssh_host.SSHHost,
41                                                     'mock SSHHost')
42        self.god.stub_function(drones.drone_utility, 'create_host')
43        self.drone_utility_path = 'mock-drone-utility-path'
44        self.mock_return = {'results': ['mock results'],
45                            'warnings': []}
46        self.god.stub_with(drones._RemoteDrone, '_drone_utility_path',
47                self.drone_utility_path)
48
49
50    def tearDown(self):
51        self.god.unstub_all()
52
53
54    def test_execute(self):
55        """Test execute method."""
56
57        drone1 = self.create_remote_drone('fakehostname1')
58        drone2 = self.create_remote_drone('fakehostname2')
59        drone3 = self.create_remote_drone('fakehostname3')
60
61        # Check task queue exception conditions.
62        task_queue = drone_task_queue.DroneTaskQueue()
63        task_queue.results[drone1] = None
64        self.assertRaises(drone_task_queue.DroneTaskQueueException,
65                          task_queue.execute, [])
66        task_queue.results.clear()
67
68        # Queue 2 calls against each drone, and confirm that the host's
69        # run method is called 3 times. Then compare results returned by
70        # the task queue against the mock results.
71        drones = [drone1, drone2, drone3]
72        for drone in drones:
73            drone.queue_call('foo')
74            drone.queue_call('bar')
75            mock_result = utils.CmdResult(
76                    stdout=cPickle.dumps(self.mock_return))
77            self._mock_host.run.expect_call(
78                    'python %s' % self.drone_utility_path,
79                    stdin=cPickle.dumps(drone.get_calls()), stdout_tee=None,
80                    connect_timeout=mock.is_instance_comparator(int)
81                    ).and_return(mock_result)
82        task_queue.execute(drones, wait=False)
83        self.assertTrue(set(task_queue.results.keys()) == set(drones))
84        results = task_queue.get_results()
85        self.assertTrue(len(task_queue.results) == 0)
86        for drone, result in results.iteritems():
87            self.assertTrue(result == self.mock_return['results'])
88
89        # Test execute and get_results
90        drone1.queue_call('foo')
91        mock_result = utils.CmdResult(stdout=cPickle.dumps(self.mock_return))
92        self._mock_host.run.expect_call(
93                'python %s' % self.drone_utility_path,
94                stdin=cPickle.dumps(drone1.get_calls()), stdout_tee=None,
95                connect_timeout=mock.is_instance_comparator(int)).and_return(
96                        mock_result)
97        self.assertTrue(task_queue.execute(drones, wait=True)[drone1] ==
98                        self.mock_return['results'])
99        self.assertTrue(len(task_queue.results) == 0)
100        self.assertTrue(len(task_queue.get_results()) == 0)
101        self.god.check_playback()
102
103
104if __name__ == '__main__':
105    unittest.main()
106