#!/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()
|
