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 16 from autotest_lib.client.common_lib.test_utils import mock 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() 106
