Home | History | Annotate | Download | only in scheduler
      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, unittest
     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