MK Ryu | 7911ad5 | 2015-12-18 11:40:04 -0800 | [diff] [blame] | 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 | |
| 9 | import cPickle |
| 10 | import logging |
| 11 | import Queue |
| 12 | import unittest |
| 13 | |
| 14 | import common |
| 15 | from autotest_lib.client.common_lib import utils |
Justin Giorgi | 67ad67d | 2016-06-29 14:41:04 -0700 | [diff] [blame] | 16 | from autotest_lib.client.common_lib.test_utils import mock |
MK Ryu | 7911ad5 | 2015-12-18 11:40:04 -0800 | [diff] [blame] | 17 | from autotest_lib.scheduler import drone_task_queue |
| 18 | from autotest_lib.scheduler import drones |
| 19 | from autotest_lib.server.hosts import ssh_host |
| 20 | |
| 21 | |
| 22 | class 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 | |
| 104 | if __name__ == '__main__': |
| 105 | unittest.main() |