| #!/usr/bin/python | |
| # | |
| # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. | |
| # Use of this source code is governed by a BSD-style license that can be | |
| # found in the LICENSE file. | |
| """Tests for the drone managers thread queue.""" | |
| import cPickle | |
| import logging | |
| import Queue | |
| import unittest | |
| import common | |
| from autotest_lib.client.common_lib import utils | |
| from autotest_lib.client.common_lib.test_utils import mock | |
| from autotest_lib.scheduler import drone_task_queue | |
| from autotest_lib.scheduler import drones | |
| from autotest_lib.scheduler import thread_lib | |
| from autotest_lib.server.hosts import ssh_host | |
| class DroneThreadLibTest(unittest.TestCase): | |
| """Threaded task queue drone library tests.""" | |
| def create_remote_drone(self, hostname): | |
| """Create and initialize a Remote Drone. | |
| @param hostname: The name of the host for the remote drone. | |
| @return: A remote drone instance. | |
| """ | |
| drones.drone_utility.create_host.expect_call(hostname).and_return( | |
| self._mock_host) | |
| self._mock_host.is_up.expect_call().and_return(True) | |
| return drones._RemoteDrone(hostname, timestamp_remote_calls=False) | |
| def setUp(self): | |
| self.god = mock.mock_god() | |
| self._mock_host = self.god.create_mock_class(ssh_host.SSHHost, | |
| 'mock SSHHost') | |
| self.god.stub_function(drones.drone_utility, 'create_host') | |
| self.drone_utility_path = 'mock-drone-utility-path' | |
| self.mock_return = {'results': ['mock results'], | |
| 'warnings': []} | |
| self.god.stub_with(drones._RemoteDrone, '_drone_utility_path', | |
| self.drone_utility_path) | |
| def tearDown(self): | |
| self.god.unstub_all() | |
| def test_worker(self): | |
| """Test the worker method of a ThreadedTaskQueue.""" | |
| # Invoke the worker method with a drone that has a queued call and check | |
| # that the drones host.run method is invoked for the call, and the | |
| # results queue contains the expected results. | |
| drone = self.create_remote_drone('fakehostname') | |
| task_queue = thread_lib.ThreadedTaskQueue() | |
| drone.queue_call('foo') | |
| mock_result = utils.CmdResult(stdout=cPickle.dumps(self.mock_return)) | |
| self._mock_host.run.expect_call( | |
| 'python %s' % self.drone_utility_path, | |
| stdin=cPickle.dumps(drone.get_calls()), stdout_tee=None, | |
| connect_timeout=mock.is_instance_comparator(int)).and_return( | |
| mock_result) | |
| task_queue.worker(drone, task_queue.results_queue) | |
| result = task_queue.results_queue.get() | |
| self.assertTrue(task_queue.results_queue.empty() and | |
| result.drone == drone and | |
| result.results == self.mock_return['results']) | |
| self.god.check_playback() | |
| def test_wait_on_drones(self): | |
| """Test waiting on drone threads.""" | |
| def waiting_func(queue): | |
| while len(queue.queue) < 2: | |
| continue | |
| logging.warning('Consuming thread finished.') | |
| queue.put(3) | |
| def exception_func(queue): | |
| while queue.empty(): | |
| continue | |
| queue.put(2) | |
| logging.warning('Failing thread raising error.') | |
| raise ValueError('Value error') | |
| def quick_func(): | |
| return | |
| # Create 2 threads, one of which raises an exception while the other | |
| # just exits normally. Insert both threads into the thread_queue against | |
| # mock drones and confirm that: | |
| # a. The task queue waits for both threads, though the first one fails. | |
| # b. The task queue records the right DroneTaskQueueException, which | |
| # contains the original exception. | |
| # c. The failing thread records its own exception instead of raising it. | |
| task_queue = thread_lib.ThreadedTaskQueue() | |
| drone1 = self.create_remote_drone('fakehostname1') | |
| drone2 = self.create_remote_drone('fakehostname2') | |
| sync_queue = Queue.Queue() | |
| waiting_worker = thread_lib.ExceptionRememberingThread( | |
| target=waiting_func, args=(sync_queue,)) | |
| failing_worker = thread_lib.ExceptionRememberingThread( | |
| target=exception_func, args=(sync_queue,)) | |
| task_queue.drone_threads[drone1] = waiting_worker | |
| task_queue.drone_threads[drone2] = failing_worker | |
| master_thread = thread_lib.ExceptionRememberingThread( | |
| target=task_queue.wait_on_drones) | |
| thread_list = [failing_worker, waiting_worker, master_thread] | |
| for thread in thread_list: | |
| thread.setDaemon(True) | |
| thread.start() | |
| sync_queue.put(1) | |
| master_thread.join() | |
| self.assertTrue(isinstance(master_thread.err, | |
| drone_task_queue.DroneTaskQueueException)) | |
| self.assertTrue(isinstance(failing_worker.err, ValueError)) | |
| self.assertTrue(str(failing_worker.err) in str(master_thread.err)) | |
| self.assertTrue(3 in list(sync_queue.queue)) | |
| self.assertTrue(task_queue.drone_threads == {}) | |
| # Call wait_on_drones after the child thread has exited. | |
| quick_worker = thread_lib.ExceptionRememberingThread(target=quick_func) | |
| task_queue.drone_threads[drone1] = quick_worker | |
| quick_worker.start() | |
| while quick_worker.isAlive(): | |
| continue | |
| task_queue.wait_on_drones() | |
| self.assertTrue(task_queue.drone_threads == {}) | |
| def test_get_results(self): | |
| """Test retrieving results from the results queue.""" | |
| # Insert results for the same drone twice into the results queue | |
| # and confirm that an exception is raised. | |
| task_queue = thread_lib.ThreadedTaskQueue() | |
| drone1 = self.create_remote_drone('fakehostname1') | |
| drone2 = self.create_remote_drone('fakehostname2') | |
| task_queue.results_queue.put( | |
| thread_lib.ThreadedTaskQueue.result(drone1, self.mock_return)) | |
| task_queue.results_queue.put( | |
| thread_lib.ThreadedTaskQueue.result(drone1, self.mock_return)) | |
| self.god.stub_function(task_queue, 'wait_on_drones') | |
| task_queue.wait_on_drones.expect_call() | |
| self.assertRaises(drone_task_queue.DroneTaskQueueException, | |
| task_queue.get_results) | |
| # Insert results for different drones and check that they're returned | |
| # in a drone results dict. | |
| self.assertTrue(task_queue.results_queue.empty()) | |
| task_queue.results_queue.put( | |
| thread_lib.ThreadedTaskQueue.result(drone1, self.mock_return)) | |
| task_queue.results_queue.put( | |
| thread_lib.ThreadedTaskQueue.result(drone2, self.mock_return)) | |
| task_queue.wait_on_drones.expect_call() | |
| results = task_queue.get_results() | |
| self.assertTrue(results[drone1] == self.mock_return and | |
| results[drone2] == self.mock_return) | |
| self.god.check_playback() | |
| def test_execute(self): | |
| """Test task queue execute.""" | |
| drone1 = self.create_remote_drone('fakehostname1') | |
| drone2 = self.create_remote_drone('fakehostname2') | |
| drone3 = self.create_remote_drone('fakehostname3') | |
| # Check task queue exception conditions. | |
| task_queue = thread_lib.ThreadedTaskQueue() | |
| task_queue.results_queue.put(1) | |
| self.assertRaises(drone_task_queue.DroneTaskQueueException, | |
| task_queue.execute, []) | |
| task_queue.results_queue.get() | |
| task_queue.drone_threads[drone1] = None | |
| self.assertRaises(drone_task_queue.DroneTaskQueueException, | |
| task_queue.execute, []) | |
| task_queue.drone_threads = {} | |
| # Queue 2 calls against each drone, and confirm that the host's | |
| # run method is called 3 times. Then check the threads created, | |
| # and finally compare results returned by the task queue against | |
| # the mock results. | |
| drones = [drone1, drone2, drone3] | |
| for drone in drones: | |
| drone.queue_call('foo') | |
| drone.queue_call('bar') | |
| mock_result = utils.CmdResult( | |
| stdout=cPickle.dumps(self.mock_return)) | |
| self._mock_host.run.expect_call( | |
| 'python %s' % self.drone_utility_path, | |
| stdin=cPickle.dumps(drone.get_calls()), stdout_tee=None, | |
| connect_timeout=mock.is_instance_comparator(int) | |
| ).and_return(mock_result) | |
| task_queue.execute(drones, wait=False) | |
| self.assertTrue(set(task_queue.drone_threads.keys()) == set(drones)) | |
| for drone, thread in task_queue.drone_threads.iteritems(): | |
| self.assertTrue(drone.hostname in thread.getName()) | |
| self.assertTrue(thread.isDaemon()) | |
| self.assertRaises(RuntimeError, thread.start) | |
| results = task_queue.get_results() | |
| for drone, result in results.iteritems(): | |
| self.assertTrue(result == self.mock_return['results']) | |
| # Test synchronous execute | |
| drone1.queue_call('foo') | |
| mock_result = utils.CmdResult(stdout=cPickle.dumps(self.mock_return)) | |
| self._mock_host.run.expect_call( | |
| 'python %s' % self.drone_utility_path, | |
| stdin=cPickle.dumps(drone1.get_calls()), stdout_tee=None, | |
| connect_timeout=mock.is_instance_comparator(int)).and_return( | |
| mock_result) | |
| self.assertTrue(task_queue.execute(drones, wait=True)[drone1] == | |
| self.mock_return['results']) | |
| self.god.check_playback() | |
| if __name__ == '__main__': | |
| unittest.main() |