Home | History | Annotate | Download | only in cros
      1 #!/usr/bin/python
      2 #
      3 # Copyright (c) 2012 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 """Unit tests for client/common_lib/cros/dev_server.py."""
      8 
      9 import __builtin__
     10 
     11 import httplib
     12 import json
     13 import mox
     14 import os
     15 import StringIO
     16 import time
     17 import unittest
     18 import urllib2
     19 
     20 import mock
     21 
     22 import common
     23 from autotest_lib.client.bin import utils as bin_utils
     24 from autotest_lib.client.common_lib import android_utils
     25 from autotest_lib.client.common_lib import error
     26 from autotest_lib.client.common_lib import global_config
     27 from autotest_lib.client.common_lib import utils
     28 from autotest_lib.client.common_lib.cros import dev_server
     29 from autotest_lib.client.common_lib.cros import retry
     30 
     31 
     32 def retry_mock(ExceptionToCheck, timeout_min, exception_to_raise=None,
     33                label=None):
     34     """A mock retry decorator to use in place of the actual one for testing.
     35 
     36     @param ExceptionToCheck: the exception to check.
     37     @param timeout_mins: Amount of time in mins to wait before timing out.
     38     @param exception_to_raise: the exception to raise in retry.retry
     39     @param label: used in debug messages
     40 
     41     """
     42     def inner_retry(func):
     43         """The actual decorator.
     44 
     45         @param func: Function to be called in decorator.
     46 
     47         """
     48         return func
     49 
     50     return inner_retry
     51 
     52 
     53 class MockSshResponse(object):
     54     """An ssh response mocked for testing."""
     55 
     56     def __init__(self, output, exit_status=0):
     57         self.stdout = output
     58         self.exit_status = exit_status
     59         self.stderr = 'SSH connection error occurred.'
     60 
     61 
     62 class MockSshError(error.CmdError):
     63     """An ssh error response mocked for testing."""
     64 
     65     def __init__(self):
     66         self.result_obj = MockSshResponse('error', exit_status=255)
     67 
     68 
     69 E403 = urllib2.HTTPError(url='',
     70                          code=httplib.FORBIDDEN,
     71                          msg='Error 403',
     72                          hdrs=None,
     73                          fp=StringIO.StringIO('Expected.'))
     74 E500 = urllib2.HTTPError(url='',
     75                          code=httplib.INTERNAL_SERVER_ERROR,
     76                          msg='Error 500',
     77                          hdrs=None,
     78                          fp=StringIO.StringIO('Expected.'))
     79 CMD_ERROR = error.CmdError('error_cmd', MockSshError().result_obj)
     80 
     81 
     82 class RunCallTest(mox.MoxTestBase):
     83     """Unit tests for ImageServerBase.run_call or DevServer.run_call."""
     84 
     85     def setUp(self):
     86         """Set up the test"""
     87         self.test_call = 'http://nothing/test'
     88         self.hostname = 'nothing'
     89         self.contents = 'true'
     90         self.contents_readline = ['file/one', 'file/two']
     91         self.save_ssh_config = dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER
     92         super(RunCallTest, self).setUp()
     93         self.mox.StubOutWithMock(urllib2, 'urlopen')
     94         self.mox.StubOutWithMock(utils, 'run')
     95 
     96         sleep = mock.patch('time.sleep', autospec=True)
     97         sleep.start()
     98         self.addCleanup(sleep.stop)
     99 
    100 
    101     def tearDown(self):
    102         """Tear down the test"""
    103         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = self.save_ssh_config
    104         super(RunCallTest, self).tearDown()
    105 
    106 
    107     def testRunCallHTTPWithDownDevserver(self):
    108         """Test dev_server.ImageServerBase.run_call using http with arg:
    109         (call)."""
    110         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = False
    111 
    112         urllib2.urlopen(mox.StrContains(self.test_call)).AndReturn(
    113                 StringIO.StringIO(dev_server.ERR_MSG_FOR_DOWN_DEVSERVER))
    114         time.sleep(mox.IgnoreArg())
    115         urllib2.urlopen(mox.StrContains(self.test_call)).AndReturn(
    116                 StringIO.StringIO(self.contents))
    117         self.mox.ReplayAll()
    118         response = dev_server.ImageServerBase.run_call(self.test_call)
    119         self.assertEquals(self.contents, response)
    120 
    121 
    122     def testRunCallSSHWithDownDevserver(self):
    123         """Test dev_server.ImageServerBase.run_call using http with arg:
    124         (call)."""
    125         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = True
    126         self.mox.StubOutWithMock(utils, 'get_restricted_subnet')
    127         utils.get_restricted_subnet(
    128                 self.hostname, utils.RESTRICTED_SUBNETS).AndReturn(
    129                 self.hostname)
    130 
    131         to_return1 = MockSshResponse(dev_server.ERR_MSG_FOR_DOWN_DEVSERVER)
    132         to_return2 = MockSshResponse(self.contents)
    133         utils.run(mox.StrContains(self.test_call),
    134                   timeout=mox.IgnoreArg()).AndReturn(to_return1)
    135         time.sleep(mox.IgnoreArg())
    136         utils.run(mox.StrContains(self.test_call),
    137                   timeout=mox.IgnoreArg()).AndReturn(to_return2)
    138 
    139         self.mox.ReplayAll()
    140         response = dev_server.ImageServerBase.run_call(self.test_call)
    141         self.assertEquals(self.contents, response)
    142         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = False
    143 
    144 
    145     def testRunCallWithSingleCallHTTP(self):
    146         """Test dev_server.ImageServerBase.run_call using http with arg:
    147         (call)."""
    148         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = False
    149 
    150         urllib2.urlopen(mox.StrContains(self.test_call)).AndReturn(
    151                 StringIO.StringIO(self.contents))
    152         self.mox.ReplayAll()
    153         response = dev_server.ImageServerBase.run_call(self.test_call)
    154         self.assertEquals(self.contents, response)
    155 
    156 
    157     def testRunCallWithCallAndReadlineHTTP(self):
    158         """Test dev_server.ImageServerBase.run_call using http with arg:
    159         (call, readline=True)."""
    160         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = False
    161 
    162         urllib2.urlopen(mox.StrContains(self.test_call)).AndReturn(
    163                 StringIO.StringIO('\n'.join(self.contents_readline)))
    164         self.mox.ReplayAll()
    165         response = dev_server.ImageServerBase.run_call(
    166                 self.test_call, readline=True)
    167         self.assertEquals(self.contents_readline, response)
    168 
    169 
    170     def testRunCallWithCallAndTimeoutHTTP(self):
    171         """Test dev_server.ImageServerBase.run_call using http with args:
    172         (call, timeout=xxx)."""
    173         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = False
    174 
    175         urllib2.urlopen(mox.StrContains(self.test_call), data=None).AndReturn(
    176                 StringIO.StringIO(self.contents))
    177         self.mox.ReplayAll()
    178         response = dev_server.ImageServerBase.run_call(
    179                 self.test_call, timeout=60)
    180         self.assertEquals(self.contents, response)
    181 
    182 
    183     def testRunCallWithSingleCallSSH(self):
    184         """Test dev_server.ImageServerBase.run_call using ssh with arg:
    185         (call)."""
    186         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = True
    187         self.mox.StubOutWithMock(utils, 'get_restricted_subnet')
    188         utils.get_restricted_subnet(
    189                 self.hostname, utils.RESTRICTED_SUBNETS).AndReturn(
    190                 self.hostname)
    191 
    192         to_return = MockSshResponse(self.contents)
    193         utils.run(mox.StrContains(self.test_call),
    194                   timeout=mox.IgnoreArg()).AndReturn(to_return)
    195         self.mox.ReplayAll()
    196         response = dev_server.ImageServerBase.run_call(self.test_call)
    197         self.assertEquals(self.contents, response)
    198 
    199 
    200     def testRunCallWithCallAndReadlineSSH(self):
    201         """Test dev_server.ImageServerBase.run_call using ssh with args:
    202         (call, readline=True)."""
    203         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = True
    204         self.mox.StubOutWithMock(utils, 'get_restricted_subnet')
    205         utils.get_restricted_subnet(
    206                 self.hostname, utils.RESTRICTED_SUBNETS).AndReturn(
    207                 self.hostname)
    208 
    209         to_return = MockSshResponse('\n'.join(self.contents_readline))
    210         utils.run(mox.StrContains(self.test_call),
    211                   timeout=mox.IgnoreArg()).AndReturn(to_return)
    212         self.mox.ReplayAll()
    213         response = dev_server.ImageServerBase.run_call(
    214                 self.test_call, readline=True)
    215         self.assertEquals(self.contents_readline, response)
    216 
    217 
    218     def testRunCallWithCallAndTimeoutSSH(self):
    219         """Test dev_server.ImageServerBase.run_call using ssh with args:
    220         (call, timeout=xxx)."""
    221         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = True
    222         self.mox.StubOutWithMock(utils, 'get_restricted_subnet')
    223         utils.get_restricted_subnet(
    224                 self.hostname, utils.RESTRICTED_SUBNETS).AndReturn(
    225                 self.hostname)
    226 
    227         to_return = MockSshResponse(self.contents)
    228         utils.run(mox.StrContains(self.test_call),
    229                   timeout=mox.IgnoreArg()).AndReturn(to_return)
    230         self.mox.ReplayAll()
    231         response = dev_server.ImageServerBase.run_call(
    232                 self.test_call, timeout=60)
    233         self.assertEquals(self.contents, response)
    234 
    235 
    236     def testRunCallWithExceptionHTTP(self):
    237         """Test dev_server.ImageServerBase.run_call using http with raising
    238         exception."""
    239         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = False
    240         urllib2.urlopen(mox.StrContains(self.test_call)).AndRaise(E500)
    241         self.mox.ReplayAll()
    242         self.assertRaises(urllib2.HTTPError,
    243                           dev_server.ImageServerBase.run_call,
    244                           self.test_call)
    245 
    246 
    247     def testRunCallWithExceptionSSH(self):
    248         """Test dev_server.ImageServerBase.run_call using ssh with raising
    249         exception."""
    250         dev_server.ENABLE_SSH_CONNECTION_FOR_DEVSERVER = True
    251         self.mox.StubOutWithMock(utils, 'get_restricted_subnet')
    252         utils.get_restricted_subnet(
    253                 self.hostname, utils.RESTRICTED_SUBNETS).AndReturn(
    254                 self.hostname)
    255 
    256         utils.run(mox.StrContains(self.test_call),
    257                   timeout=mox.IgnoreArg()).AndRaise(MockSshError())
    258         self.mox.ReplayAll()
    259         self.assertRaises(error.CmdError,
    260                           dev_server.ImageServerBase.run_call,
    261                           self.test_call)
    262 
    263 
    264     def testRunCallByDevServerHTTP(self):
    265         """Test dev_server.DevServer.run_call, which uses http, and can be
    266         directly called by CrashServer."""
    267         urllib2.urlopen(
    268                 mox.StrContains(self.test_call), data=None).AndReturn(
    269                         StringIO.StringIO(self.contents))
    270         self.mox.ReplayAll()
    271         response = dev_server.DevServer.run_call(
    272                self.test_call, timeout=60)
    273         self.assertEquals(self.contents, response)
    274 
    275 
    276 class DevServerTest(mox.MoxTestBase):
    277     """Unit tests for dev_server.DevServer.
    278 
    279     @var _HOST: fake dev server host address.
    280     """
    281 
    282     _HOST = 'http://nothing'
    283     _CRASH_HOST = 'http://nothing-crashed'
    284     _CONFIG = global_config.global_config
    285 
    286 
    287     def setUp(self):
    288         """Set up the test"""
    289         super(DevServerTest, self).setUp()
    290         self.crash_server = dev_server.CrashServer(DevServerTest._CRASH_HOST)
    291         self.dev_server = dev_server.ImageServer(DevServerTest._HOST)
    292         self.android_dev_server = dev_server.AndroidBuildServer(
    293                 DevServerTest._HOST)
    294         self.mox.StubOutWithMock(dev_server.ImageServerBase, 'run_call')
    295         self.mox.StubOutWithMock(urllib2, 'urlopen')
    296         self.mox.StubOutWithMock(utils, 'run')
    297         self.mox.StubOutWithMock(os.path, 'exists')
    298         # Hide local restricted_subnets setting.
    299         dev_server.RESTRICTED_SUBNETS = []
    300         self.mox.StubOutWithMock(dev_server.ImageServer,
    301                                  '_read_json_response_from_devserver')
    302 
    303         sleep = mock.patch('time.sleep', autospec=True)
    304         sleep.start()
    305         self.addCleanup(sleep.stop)
    306 
    307 
    308     def testSimpleResolve(self):
    309         """One devserver, verify we resolve to it."""
    310         self.mox.StubOutWithMock(dev_server, '_get_dev_server_list')
    311         self.mox.StubOutWithMock(dev_server.ImageServer, 'devserver_healthy')
    312         dev_server._get_dev_server_list().MultipleTimes().AndReturn(
    313                 [DevServerTest._HOST])
    314         dev_server.ImageServer.devserver_healthy(DevServerTest._HOST).AndReturn(
    315                                                                         True)
    316         self.mox.ReplayAll()
    317         devserver = dev_server.ImageServer.resolve('my_build')
    318         self.assertEquals(devserver.url(), DevServerTest._HOST)
    319 
    320 
    321     def testResolveWithFailure(self):
    322         """Ensure we rehash on a failed ping on a bad_host."""
    323         self.mox.StubOutWithMock(dev_server, '_get_dev_server_list')
    324         bad_host, good_host = 'http://bad_host:99', 'http://good_host:8080'
    325         dev_server._get_dev_server_list().MultipleTimes().AndReturn(
    326                 [bad_host, good_host])
    327         argument1 = mox.StrContains(bad_host)
    328         argument2 = mox.StrContains(good_host)
    329 
    330         # Mock out bad ping failure to bad_host by raising devserver exception.
    331         dev_server.ImageServerBase.run_call(
    332                 argument1, timeout=mox.IgnoreArg()).AndRaise(
    333                         dev_server.DevServerException())
    334         # Good host is good.
    335         dev_server.ImageServerBase.run_call(
    336                 argument2, timeout=mox.IgnoreArg()).AndReturn(
    337                         '{"free_disk": 1024}')
    338 
    339         self.mox.ReplayAll()
    340         host = dev_server.ImageServer.resolve(0) # Using 0 as it'll hash to 0.
    341         self.assertEquals(host.url(), good_host)
    342         self.mox.VerifyAll()
    343 
    344 
    345     def testResolveWithFailureURLError(self):
    346         """Ensure we rehash on a failed ping using http on a bad_host after
    347         urlerror."""
    348         # Set retry.retry to retry_mock for just returning the original
    349         # method for this test. This is to save waiting time for real retry,
    350         # which is defined by dev_server.DEVSERVER_SSH_TIMEOUT_MINS.
    351         # Will reset retry.retry to real retry at the end of this test.
    352         real_retry = retry.retry
    353         retry.retry = retry_mock
    354 
    355         self.mox.StubOutWithMock(dev_server, '_get_dev_server_list')
    356         bad_host, good_host = 'http://bad_host:99', 'http://good_host:8080'
    357         dev_server._get_dev_server_list().MultipleTimes().AndReturn(
    358                 [bad_host, good_host])
    359         argument1 = mox.StrContains(bad_host)
    360         argument2 = mox.StrContains(good_host)
    361 
    362         # Mock out bad ping failure to bad_host by raising devserver exception.
    363         dev_server.ImageServerBase.run_call(
    364                 argument1, timeout=mox.IgnoreArg()).MultipleTimes().AndRaise(
    365                         urllib2.URLError('urlopen connection timeout'))
    366 
    367         # Good host is good.
    368         dev_server.ImageServerBase.run_call(
    369                 argument2, timeout=mox.IgnoreArg()).AndReturn(
    370                         '{"free_disk": 1024}')
    371 
    372         self.mox.ReplayAll()
    373         host = dev_server.ImageServer.resolve(0) # Using 0 as it'll hash to 0.
    374         self.assertEquals(host.url(), good_host)
    375         self.mox.VerifyAll()
    376 
    377         retry.retry = real_retry
    378 
    379 
    380     def testResolveWithManyDevservers(self):
    381         """Should be able to return different urls with multiple devservers."""
    382         self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
    383         self.mox.StubOutWithMock(dev_server.DevServer, 'devserver_healthy')
    384 
    385         host0_expected = 'http://host0:8080'
    386         host1_expected = 'http://host1:8082'
    387 
    388         dev_server.ImageServer.servers().MultipleTimes().AndReturn(
    389                 [host0_expected, host1_expected])
    390         dev_server.ImageServer.devserver_healthy(host0_expected).AndReturn(True)
    391         dev_server.ImageServer.devserver_healthy(host1_expected).AndReturn(True)
    392 
    393         self.mox.ReplayAll()
    394         host0 = dev_server.ImageServer.resolve(0)
    395         host1 = dev_server.ImageServer.resolve(1)
    396         self.mox.VerifyAll()
    397 
    398         self.assertEqual(host0.url(), host0_expected)
    399         self.assertEqual(host1.url(), host1_expected)
    400 
    401 
    402     def testCmdErrorRetryCollectAULog(self):
    403         """Devserver should retry _collect_au_log() on CMDError,
    404         but pass through real exception."""
    405         dev_server.ImageServerBase.run_call(
    406                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
    407         dev_server.ImageServerBase.run_call(
    408                 mox.IgnoreArg()).AndRaise(E500)
    409         self.mox.ReplayAll()
    410         self.assertFalse(self.dev_server.collect_au_log(
    411                 '100.0.0.0', 100, 'path/'))
    412 
    413 
    414     def testURLErrorRetryCollectAULog(self):
    415         """Devserver should retry _collect_au_log() on URLError,
    416         but pass through real exception."""
    417         self.mox.StubOutWithMock(time, 'sleep')
    418 
    419         refused = urllib2.URLError('[Errno 111] Connection refused')
    420         dev_server.ImageServerBase.run_call(
    421                 mox.IgnoreArg()).AndRaise(refused)
    422         time.sleep(mox.IgnoreArg())
    423         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E403)
    424         self.mox.ReplayAll()
    425         self.assertFalse(self.dev_server.collect_au_log(
    426                 '100.0.0.0', 100, 'path/'))
    427 
    428 
    429     def testCmdErrorRetryKillAUProcess(self):
    430         """Devserver should retry _kill_au_process() on CMDError,
    431         but pass through real exception."""
    432         dev_server.ImageServerBase.run_call(
    433                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
    434         dev_server.ImageServerBase.run_call(
    435                 mox.IgnoreArg()).AndRaise(E500)
    436         self.mox.ReplayAll()
    437         self.assertFalse(self.dev_server.kill_au_process_for_host(
    438                 '100.0.0.0', 100))
    439 
    440 
    441     def testURLErrorRetryKillAUProcess(self):
    442         """Devserver should retry _kill_au_process() on URLError,
    443         but pass through real exception."""
    444         self.mox.StubOutWithMock(time, 'sleep')
    445 
    446         refused = urllib2.URLError('[Errno 111] Connection refused')
    447         dev_server.ImageServerBase.run_call(
    448                 mox.IgnoreArg()).AndRaise(refused)
    449         time.sleep(mox.IgnoreArg())
    450         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E403)
    451         self.mox.ReplayAll()
    452         self.assertFalse(self.dev_server.kill_au_process_for_host(
    453                 '100.0.0.0', 100))
    454 
    455 
    456     def testCmdErrorRetryCleanTrackLog(self):
    457         """Devserver should retry _clean_track_log() on CMDError,
    458         but pass through real exception."""
    459         dev_server.ImageServerBase.run_call(
    460                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
    461         dev_server.ImageServerBase.run_call(
    462                 mox.IgnoreArg()).AndRaise(E500)
    463         self.mox.ReplayAll()
    464         self.assertFalse(self.dev_server.clean_track_log('100.0.0.0', 100))
    465 
    466 
    467     def testURLErrorRetryCleanTrackLog(self):
    468         """Devserver should retry _clean_track_log() on URLError,
    469         but pass through real exception."""
    470         self.mox.StubOutWithMock(time, 'sleep')
    471 
    472         refused = urllib2.URLError('[Errno 111] Connection refused')
    473         dev_server.ImageServerBase.run_call(
    474                 mox.IgnoreArg()).AndRaise(refused)
    475         time.sleep(mox.IgnoreArg())
    476         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E403)
    477         self.mox.ReplayAll()
    478         self.assertFalse(self.dev_server.clean_track_log('100.0.0.0', 100))
    479 
    480 
    481     def _mockWriteFile(self):
    482         """Mock write content to a file."""
    483         mock_file = self.mox.CreateMockAnything()
    484         open(mox.IgnoreArg(), 'w').AndReturn(mock_file)
    485         mock_file.__enter__().AndReturn(mock_file)
    486         mock_file.write(mox.IgnoreArg())
    487         mock_file.__exit__(None, None, None)
    488 
    489 
    490     def _preSetupForAutoUpdate(self, **kwargs):
    491         """Pre-setup for testing auto_update logics and error handling in
    492         devserver."""
    493         response1 = (True, 100)
    494         response2 = (True, 'Completed')
    495         response3 = {'host_logs': {'a': 'log'}, 'cros_au_log': 'logs'}
    496 
    497         argument1 = mox.And(mox.StrContains(self._HOST),
    498                             mox.StrContains('cros_au'))
    499         argument2 = mox.And(mox.StrContains(self._HOST),
    500                             mox.StrContains('get_au_status'))
    501         argument3 = mox.And(mox.StrContains(self._HOST),
    502                             mox.StrContains('collect_cros_au_log'))
    503         argument4 = mox.And(mox.StrContains(self._HOST),
    504                             mox.StrContains('handler_cleanup'))
    505         argument5 = mox.And(mox.StrContains(self._HOST),
    506                             mox.StrContains('kill_au_proc'))
    507 
    508         retry_error = None
    509         if 'retry_error' in kwargs:
    510             retry_error = kwargs['retry_error']
    511 
    512         raised_error = E403
    513         if 'raised_error' in kwargs:
    514             raised_error = kwargs['raised_error']
    515 
    516         if 'cros_au_error' in kwargs:
    517             if retry_error:
    518                 dev_server.ImageServerBase.run_call(argument1).AndRaise(
    519                         retry_error)
    520                 time.sleep(mox.IgnoreArg())
    521 
    522             if kwargs['cros_au_error']:
    523                 dev_server.ImageServerBase.run_call(argument1).AndRaise(
    524                         raised_error)
    525             else:
    526                 dev_server.ImageServerBase.run_call(argument1).AndReturn(
    527                         json.dumps(response1))
    528 
    529         if 'get_au_status_error' in kwargs:
    530             if retry_error:
    531                 dev_server.ImageServerBase.run_call(argument2).AndRaise(
    532                         retry_error)
    533                 time.sleep(mox.IgnoreArg())
    534 
    535             if kwargs['get_au_status_error']:
    536                 dev_server.ImageServerBase.run_call(argument2).AndRaise(
    537                         raised_error)
    538             else:
    539                 dev_server.ImageServerBase.run_call(argument2).AndReturn(
    540                         json.dumps(response2))
    541 
    542         if 'collect_au_log_error' in kwargs:
    543             if kwargs['collect_au_log_error']:
    544                 dev_server.ImageServerBase.run_call(argument3).AndRaise(
    545                         raised_error)
    546             else:
    547                 dev_server.ImageServer._read_json_response_from_devserver(
    548                         mox.IgnoreArg()).AndReturn(response3)
    549                 dev_server.ImageServerBase.run_call(argument3).AndReturn('log')
    550                 os.path.exists(mox.IgnoreArg()).AndReturn(True)
    551 
    552                 # We write two log files: host_log and cros_au_log
    553                 self._mockWriteFile()
    554                 self._mockWriteFile()
    555 
    556         if 'handler_cleanup_error' in kwargs:
    557             if kwargs['handler_cleanup_error']:
    558                 dev_server.ImageServerBase.run_call(argument4).AndRaise(
    559                         raised_error)
    560             else:
    561                 dev_server.ImageServerBase.run_call(argument4).AndReturn('True')
    562 
    563         if 'kill_au_proc_error' in kwargs:
    564             if kwargs['kill_au_proc_error']:
    565                 dev_server.ImageServerBase.run_call(argument5).AndRaise(
    566                         raised_error)
    567             else:
    568                 dev_server.ImageServerBase.run_call(argument5).AndReturn('True')
    569 
    570 
    571     def testSuccessfulTriggerAutoUpdate(self):
    572         """Verify the dev server's auto_update() succeeds."""
    573         kwargs={'cros_au_error': False, 'get_au_status_error': False,
    574                 'handler_cleanup_error': False}
    575         self._preSetupForAutoUpdate(**kwargs)
    576 
    577         self.mox.ReplayAll()
    578         self.dev_server.auto_update('100.0.0.0', '')
    579         self.mox.VerifyAll()
    580 
    581 
    582     def testSuccessfulTriggerAutoUpdateWithCollectingLog(self):
    583         """Verify the dev server's auto_update() with collecting logs
    584         succeeds."""
    585         kwargs={'cros_au_error': False, 'get_au_status_error': False,
    586                 'handler_cleanup_error': False, 'collect_au_log_error': False}
    587         self.mox.StubOutWithMock(__builtin__, 'open')
    588         self._preSetupForAutoUpdate(**kwargs)
    589 
    590         self.mox.ReplayAll()
    591         self.dev_server.auto_update('100.0.0.0', '', log_dir='path/')
    592         self.mox.VerifyAll()
    593 
    594 
    595     def testCrOSAUURLErrorRetryTriggerAutoUpdateSucceed(self):
    596         """Devserver should retry cros_au() on URLError."""
    597         self.mox.StubOutWithMock(time, 'sleep')
    598         refused = urllib2.URLError('[Errno 111] Connection refused')
    599         kwargs={'retry_error': refused, 'cros_au_error': False,
    600                 'get_au_status_error': False, 'handler_cleanup_error': False,
    601                 'collect_au_log_error': False}
    602         self.mox.StubOutWithMock(__builtin__, 'open')
    603         self._preSetupForAutoUpdate(**kwargs)
    604 
    605         self.mox.ReplayAll()
    606         self.dev_server.auto_update('100.0.0.0', '', log_dir='path/')
    607         self.mox.VerifyAll()
    608 
    609 
    610     def testCrOSAUCmdErrorRetryTriggerAutoUpdateSucceed(self):
    611         """Devserver should retry cros_au() on CMDError."""
    612         self.mox.StubOutWithMock(time, 'sleep')
    613         self.mox.StubOutWithMock(__builtin__, 'open')
    614         kwargs={'retry_error': CMD_ERROR, 'cros_au_error': False,
    615                 'get_au_status_error': False, 'handler_cleanup_error': False,
    616                 'collect_au_log_error': False}
    617         self._preSetupForAutoUpdate(**kwargs)
    618 
    619         self.mox.ReplayAll()
    620         self.dev_server.auto_update('100.0.0.0', '', log_dir='path/')
    621         self.mox.VerifyAll()
    622 
    623 
    624     def testCrOSAUURLErrorRetryTriggerAutoUpdateFail(self):
    625         """Devserver should retry cros_au() on URLError, but pass through
    626         real exception."""
    627         self.mox.StubOutWithMock(time, 'sleep')
    628         refused = urllib2.URLError('[Errno 111] Connection refused')
    629         kwargs={'retry_error': refused, 'cros_au_error': True,
    630                 'raised_error': E500}
    631 
    632         for i in range(dev_server.AU_RETRY_LIMIT):
    633             self._preSetupForAutoUpdate(**kwargs)
    634             if i < dev_server.AU_RETRY_LIMIT - 1:
    635                 time.sleep(mox.IgnoreArg())
    636 
    637         self.mox.ReplayAll()
    638         self.assertRaises(dev_server.DevServerException,
    639                           self.dev_server.auto_update,
    640                           '', '')
    641 
    642 
    643     def testCrOSAUCmdErrorRetryTriggerAutoUpdateFail(self):
    644         """Devserver should retry cros_au() on CMDError, but pass through
    645         real exception."""
    646         self.mox.StubOutWithMock(time, 'sleep')
    647         kwargs={'retry_error': CMD_ERROR, 'cros_au_error': True}
    648 
    649         for i in range(dev_server.AU_RETRY_LIMIT):
    650             self._preSetupForAutoUpdate(**kwargs)
    651             if i < dev_server.AU_RETRY_LIMIT - 1:
    652                 time.sleep(mox.IgnoreArg())
    653 
    654         self.mox.ReplayAll()
    655         self.assertRaises(dev_server.DevServerException,
    656                           self.dev_server.auto_update,
    657                           '', '')
    658 
    659     def testBadBuildInAutoUpdate(self):
    660         """Devsever raises BadBuildException when DUT raises RootfsUpdateError.
    661         """
    662         kwargs = {
    663             'cros_au_error': False, 'get_au_status_error': True,
    664             'handler_cleanup_error': False,
    665             'kill_au_proc_error': False,
    666             'raised_error': urllib2.HTTPError(
    667                 code=httplib.INTERNAL_SERVER_ERROR, msg='', url='', hdrs=None,
    668                 fp=StringIO.StringIO(
    669                     'RootfsUpdateError: Build xyz failed to boot on dut.'))
    670         }
    671         for _ in range(dev_server.AU_RETRY_LIMIT):
    672             self._preSetupForAutoUpdate(**kwargs)
    673         self.mox.ReplayAll()
    674         self.assertRaises(
    675             dev_server.BadBuildException,
    676             self.dev_server.auto_update, '100.0.0.0', ''
    677         )
    678 
    679     def testGetAUStatusErrorInAutoUpdate(self):
    680         """Verify devserver's auto_update() logics for handling get_au_status
    681         errors.
    682 
    683         Func auto_update() should call 'handler_cleanup' and 'collect_au_log'
    684         even if '_trigger_auto_update()' failed.
    685         """
    686         self.mox.StubOutWithMock(time, 'sleep')
    687         self.mox.StubOutWithMock(__builtin__, 'open')
    688         kwargs={'cros_au_error': False, 'get_au_status_error': True,
    689                 'handler_cleanup_error': False, 'collect_au_log_error': False,
    690                 'kill_au_proc_error': False}
    691 
    692         for i in range(dev_server.AU_RETRY_LIMIT):
    693             self._preSetupForAutoUpdate(**kwargs)
    694             if i < dev_server.AU_RETRY_LIMIT - 1:
    695                 time.sleep(mox.IgnoreArg())
    696 
    697         self.mox.ReplayAll()
    698         self.assertRaises(dev_server.DevServerException,
    699                           self.dev_server.auto_update,
    700                           '100.0.0.0', 'build', log_dir='path/')
    701 
    702 
    703     def testCleanUpErrorInAutoUpdate(self):
    704         """Verify devserver's auto_update() logics for handling handler_cleanup
    705         errors.
    706 
    707         Func auto_update() should call 'handler_cleanup' and 'collect_au_log'
    708         no matter '_trigger_auto_update()' succeeds or fails.
    709         """
    710         self.mox.StubOutWithMock(time, 'sleep')
    711         self.mox.StubOutWithMock(__builtin__, 'open')
    712         kwargs={'cros_au_error': False, 'get_au_status_error': False,
    713                 'handler_cleanup_error': True, 'collect_au_log_error': False,
    714                 'kill_au_proc_error': False}
    715 
    716 
    717         for i in range(dev_server.AU_RETRY_LIMIT):
    718             self._preSetupForAutoUpdate(**kwargs)
    719             if i < dev_server.AU_RETRY_LIMIT - 1:
    720                 time.sleep(mox.IgnoreArg())
    721 
    722         self.mox.ReplayAll()
    723         self.assertRaises(dev_server.DevServerException,
    724                           self.dev_server.auto_update,
    725                           '100.0.0.0', 'build', log_dir='path/')
    726 
    727 
    728     def testCollectLogErrorInAutoUpdate(self):
    729         """Verify devserver's auto_update() logics for handling collect_au_log
    730         errors."""
    731         self.mox.StubOutWithMock(time, 'sleep')
    732         kwargs={'cros_au_error': False, 'get_au_status_error': False,
    733                 'handler_cleanup_error': False, 'collect_au_log_error': True,
    734                 'kill_au_proc_error': False}
    735 
    736 
    737         for i in range(dev_server.AU_RETRY_LIMIT):
    738             self._preSetupForAutoUpdate(**kwargs)
    739             if i < dev_server.AU_RETRY_LIMIT - 1:
    740                 time.sleep(mox.IgnoreArg())
    741 
    742         self.mox.ReplayAll()
    743         self.assertRaises(dev_server.DevServerException,
    744                           self.dev_server.auto_update,
    745                           '100.0.0.0', 'build', log_dir='path/')
    746 
    747 
    748     def testGetAUStatusErrorAndCleanUpErrorInAutoUpdate(self):
    749         """Verify devserver's auto_update() logics for handling get_au_status
    750         and handler_cleanup errors.
    751 
    752         Func auto_update() should call 'handler_cleanup' and 'collect_au_log'
    753         even if '_trigger_auto_update()' fails.
    754         """
    755         self.mox.StubOutWithMock(time, 'sleep')
    756         self.mox.StubOutWithMock(__builtin__, 'open')
    757         kwargs={'cros_au_error': False, 'get_au_status_error': True,
    758                 'handler_cleanup_error': True, 'collect_au_log_error': False,
    759                 'kill_au_proc_error': False}
    760 
    761 
    762         for i in range(dev_server.AU_RETRY_LIMIT):
    763             self._preSetupForAutoUpdate(**kwargs)
    764             if i < dev_server.AU_RETRY_LIMIT - 1:
    765                 time.sleep(mox.IgnoreArg())
    766 
    767         self.mox.ReplayAll()
    768         self.assertRaises(dev_server.DevServerException,
    769                           self.dev_server.auto_update,
    770                           '100.0.0.0', 'build', log_dir='path/')
    771 
    772 
    773     def testGetAUStatusErrorAndCleanUpErrorAndCollectLogErrorInAutoUpdate(self):
    774         """Verify devserver's auto_update() logics for handling get_au_status,
    775         handler_cleanup, and collect_au_log errors.
    776 
    777         Func auto_update() should call 'handler_cleanup' and 'collect_au_log'
    778         even if '_trigger_auto_update()' fails.
    779         """
    780         self.mox.StubOutWithMock(time, 'sleep')
    781         kwargs={'cros_au_error': False, 'get_au_status_error': True,
    782                 'handler_cleanup_error': True, 'collect_au_log_error': True,
    783                 'kill_au_proc_error': False}
    784 
    785         for i in range(dev_server.AU_RETRY_LIMIT):
    786             self._preSetupForAutoUpdate(**kwargs)
    787             if i < dev_server.AU_RETRY_LIMIT - 1:
    788                 time.sleep(mox.IgnoreArg())
    789 
    790         self.mox.ReplayAll()
    791         self.assertRaises(dev_server.DevServerException,
    792                           self.dev_server.auto_update,
    793                           '100.0.0.0', 'build', log_dir='path/')
    794 
    795 
    796     def testGetAUStatusErrorAndCleanUpErrorAndCollectLogErrorAndKillErrorInAutoUpdate(self):
    797         """Verify devserver's auto_update() logics for handling get_au_status,
    798         handler_cleanup, collect_au_log, and kill_au_proc errors.
    799 
    800         Func auto_update() should call 'handler_cleanup' and 'collect_au_log'
    801         even if '_trigger_auto_update()' fails.
    802         """
    803         self.mox.StubOutWithMock(time, 'sleep')
    804 
    805         kwargs={'cros_au_error': False, 'get_au_status_error': True,
    806                 'handler_cleanup_error': True, 'collect_au_log_error': True,
    807                 'kill_au_proc_error': True}
    808 
    809         for i in range(dev_server.AU_RETRY_LIMIT):
    810             self._preSetupForAutoUpdate(**kwargs)
    811             if i < dev_server.AU_RETRY_LIMIT - 1:
    812                 time.sleep(mox.IgnoreArg())
    813 
    814         self.mox.ReplayAll()
    815         self.assertRaises(dev_server.DevServerException,
    816                           self.dev_server.auto_update,
    817                           '100.0.0.0', 'build', log_dir='path/')
    818 
    819 
    820     def testSuccessfulTriggerDownloadSync(self):
    821         """Call the dev server's download method with synchronous=True."""
    822         name = 'fake/image'
    823         self.mox.StubOutWithMock(dev_server.ImageServer, '_finish_download')
    824         argument1 = mox.And(mox.StrContains(self._HOST), mox.StrContains(name),
    825                             mox.StrContains('stage?'))
    826         argument2 = mox.And(mox.StrContains(self._HOST), mox.StrContains(name),
    827                             mox.StrContains('is_staged'))
    828         dev_server.ImageServerBase.run_call(argument1).AndReturn('Success')
    829         dev_server.ImageServerBase.run_call(argument2).AndReturn('True')
    830         self.dev_server._finish_download(name, mox.IgnoreArg(), mox.IgnoreArg())
    831 
    832         # Synchronous case requires a call to finish download.
    833         self.mox.ReplayAll()
    834         self.dev_server.trigger_download(name, synchronous=True)
    835         self.mox.VerifyAll()
    836 
    837 
    838     def testSuccessfulTriggerDownloadASync(self):
    839         """Call the dev server's download method with synchronous=False."""
    840         name = 'fake/image'
    841         argument1 = mox.And(mox.StrContains(self._HOST), mox.StrContains(name),
    842                             mox.StrContains('stage?'))
    843         argument2 = mox.And(mox.StrContains(self._HOST), mox.StrContains(name),
    844                             mox.StrContains('is_staged'))
    845         dev_server.ImageServerBase.run_call(argument1).AndReturn('Success')
    846         dev_server.ImageServerBase.run_call(argument2).AndReturn('True')
    847 
    848         self.mox.ReplayAll()
    849         self.dev_server.trigger_download(name, synchronous=False)
    850         self.mox.VerifyAll()
    851 
    852 
    853     def testURLErrorRetryTriggerDownload(self):
    854         """Should retry on URLError, but pass through real exception."""
    855         self.mox.StubOutWithMock(time, 'sleep')
    856 
    857         refused = urllib2.URLError('[Errno 111] Connection refused')
    858         dev_server.ImageServerBase.run_call(
    859                 mox.IgnoreArg()).AndRaise(refused)
    860         time.sleep(mox.IgnoreArg())
    861         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E403)
    862         self.mox.ReplayAll()
    863         self.assertRaises(dev_server.DevServerException,
    864                           self.dev_server.trigger_download,
    865                           '')
    866 
    867 
    868     def testErrorTriggerDownload(self):
    869         """Should call the dev server's download method using http, fail
    870         gracefully."""
    871         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E500)
    872         self.mox.ReplayAll()
    873         self.assertRaises(dev_server.DevServerException,
    874                           self.dev_server.trigger_download,
    875                           '')
    876 
    877 
    878     def testForbiddenTriggerDownload(self):
    879         """Should call the dev server's download method using http,
    880         get exception."""
    881         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E403)
    882         self.mox.ReplayAll()
    883         self.assertRaises(dev_server.DevServerException,
    884                           self.dev_server.trigger_download,
    885                           '')
    886 
    887 
    888     def testCmdErrorTriggerDownload(self):
    889         """Should call the dev server's download method using ssh, retry
    890         trigger_download when getting error.CmdError, raise exception for
    891         urllib2.HTTPError."""
    892         dev_server.ImageServerBase.run_call(
    893                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
    894         dev_server.ImageServerBase.run_call(
    895                 mox.IgnoreArg()).AndRaise(E500)
    896         self.mox.ReplayAll()
    897         self.assertRaises(dev_server.DevServerException,
    898                           self.dev_server.trigger_download,
    899                           '')
    900 
    901 
    902     def testSuccessfulFinishDownload(self):
    903         """Should successfully call the dev server's finish download method."""
    904         name = 'fake/image'
    905         argument1 = mox.And(mox.StrContains(self._HOST),
    906                             mox.StrContains(name),
    907                             mox.StrContains('stage?'))
    908         argument2 = mox.And(mox.StrContains(self._HOST),
    909                             mox.StrContains(name),
    910                             mox.StrContains('is_staged'))
    911         dev_server.ImageServerBase.run_call(argument1).AndReturn('Success')
    912         dev_server.ImageServerBase.run_call(argument2).AndReturn('True')
    913 
    914         # Synchronous case requires a call to finish download.
    915         self.mox.ReplayAll()
    916         self.dev_server.finish_download(name)  # Raises on failure.
    917         self.mox.VerifyAll()
    918 
    919 
    920     def testErrorFinishDownload(self):
    921         """Should call the dev server's finish download method using http, fail
    922         gracefully."""
    923         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E500)
    924         self.mox.ReplayAll()
    925         self.assertRaises(dev_server.DevServerException,
    926                           self.dev_server.finish_download,
    927                           '')
    928 
    929 
    930     def testCmdErrorFinishDownload(self):
    931         """Should call the dev server's finish download method using ssh,
    932         retry finish_download when getting error.CmdError, raise exception
    933         for urllib2.HTTPError."""
    934         dev_server.ImageServerBase.run_call(
    935                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
    936         dev_server.ImageServerBase.run_call(
    937                 mox.IgnoreArg()).AndRaise(E500)
    938         self.mox.ReplayAll()
    939         self.assertRaises(dev_server.DevServerException,
    940                           self.dev_server.finish_download,
    941                           '')
    942 
    943 
    944     def testListControlFiles(self):
    945         """Should successfully list control files from the dev server."""
    946         name = 'fake/build'
    947         control_files = ['file/one', 'file/two']
    948         argument = mox.And(mox.StrContains(self._HOST),
    949                            mox.StrContains(name))
    950         dev_server.ImageServerBase.run_call(
    951                 argument, readline=True).AndReturn(control_files)
    952 
    953         self.mox.ReplayAll()
    954         paths = self.dev_server.list_control_files(name)
    955         self.assertEquals(len(paths), 2)
    956         for f in control_files:
    957             self.assertTrue(f in paths)
    958 
    959 
    960     def testFailedListControlFiles(self):
    961         """Should call the dev server's list-files method using http, get
    962         exception."""
    963         dev_server.ImageServerBase.run_call(
    964                 mox.IgnoreArg(), readline=True).AndRaise(E500)
    965         self.mox.ReplayAll()
    966         self.assertRaises(dev_server.DevServerException,
    967                           self.dev_server.list_control_files,
    968                           '')
    969 
    970 
    971     def testExplodingListControlFiles(self):
    972         """Should call the dev server's list-files method using http, get
    973         exception."""
    974         dev_server.ImageServerBase.run_call(
    975                 mox.IgnoreArg(), readline=True).AndRaise(E403)
    976         self.mox.ReplayAll()
    977         self.assertRaises(dev_server.DevServerException,
    978                           self.dev_server.list_control_files,
    979                           '')
    980 
    981 
    982     def testCmdErrorListControlFiles(self):
    983         """Should call the dev server's list-files method using ssh, retry
    984         list_control_files when getting error.CmdError, raise exception for
    985         urllib2.HTTPError."""
    986         dev_server.ImageServerBase.run_call(
    987                 mox.IgnoreArg(), readline=True).AndRaise(CMD_ERROR)
    988         dev_server.ImageServerBase.run_call(
    989                 mox.IgnoreArg(), readline=True).AndRaise(E500)
    990         self.mox.ReplayAll()
    991         self.assertRaises(dev_server.DevServerException,
    992                           self.dev_server.list_control_files,
    993                           '')
    994 
    995     def testListSuiteControls(self):
    996         """Should successfully list all contents of control files from the dev
    997         server."""
    998         name = 'fake/build'
    999         control_contents = ['control file one', 'control file two']
   1000         argument = mox.And(mox.StrContains(self._HOST),
   1001                            mox.StrContains(name))
   1002         dev_server.ImageServerBase.run_call(
   1003                 argument).AndReturn(json.dumps(control_contents))
   1004 
   1005         self.mox.ReplayAll()
   1006         file_contents = self.dev_server.list_suite_controls(name)
   1007         self.assertEquals(len(file_contents), 2)
   1008         for f in control_contents:
   1009             self.assertTrue(f in file_contents)
   1010 
   1011 
   1012     def testFailedListSuiteControls(self):
   1013         """Should call the dev server's list_suite_controls method using http,
   1014         get exception."""
   1015         dev_server.ImageServerBase.run_call(
   1016                 mox.IgnoreArg()).AndRaise(E500)
   1017         self.mox.ReplayAll()
   1018         self.assertRaises(dev_server.DevServerException,
   1019                           self.dev_server.list_suite_controls,
   1020                           '')
   1021 
   1022 
   1023     def testExplodingListSuiteControls(self):
   1024         """Should call the dev server's list_suite_controls method using http,
   1025         get exception."""
   1026         dev_server.ImageServerBase.run_call(
   1027                 mox.IgnoreArg()).AndRaise(E403)
   1028         self.mox.ReplayAll()
   1029         self.assertRaises(dev_server.DevServerException,
   1030                           self.dev_server.list_suite_controls,
   1031                           '')
   1032 
   1033 
   1034     def testCmdErrorListSuiteControls(self):
   1035         """Should call the dev server's list_suite_controls method using ssh,
   1036         retry list_suite_controls when getting error.CmdError, raise exception
   1037         for urllib2.HTTPError."""
   1038         dev_server.ImageServerBase.run_call(
   1039                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
   1040         dev_server.ImageServerBase.run_call(
   1041                 mox.IgnoreArg()).AndRaise(E500)
   1042         self.mox.ReplayAll()
   1043         self.assertRaises(dev_server.DevServerException,
   1044                           self.dev_server.list_suite_controls,
   1045                           '')
   1046 
   1047 
   1048     def testGetControlFile(self):
   1049         """Should successfully get a control file from the dev server."""
   1050         name = 'fake/build'
   1051         file = 'file/one'
   1052         contents = 'Multi-line\nControl File Contents\n'
   1053         argument = mox.And(mox.StrContains(self._HOST),
   1054                             mox.StrContains(name),
   1055                             mox.StrContains(file))
   1056         dev_server.ImageServerBase.run_call(argument).AndReturn(contents)
   1057 
   1058         self.mox.ReplayAll()
   1059         self.assertEquals(self.dev_server.get_control_file(name, file),
   1060                           contents)
   1061 
   1062 
   1063     def testErrorGetControlFile(self):
   1064         """Should try to get the contents of a control file using http, get
   1065         exception."""
   1066         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E500)
   1067         self.mox.ReplayAll()
   1068         self.assertRaises(dev_server.DevServerException,
   1069                           self.dev_server.get_control_file,
   1070                           '', '')
   1071 
   1072 
   1073     def testForbiddenGetControlFile(self):
   1074         """Should try to get the contents of a control file using http, get
   1075         exception."""
   1076         dev_server.ImageServerBase.run_call(mox.IgnoreArg()).AndRaise(E403)
   1077         self.mox.ReplayAll()
   1078         self.assertRaises(dev_server.DevServerException,
   1079                           self.dev_server.get_control_file,
   1080                           '', '')
   1081 
   1082 
   1083     def testCmdErrorGetControlFile(self):
   1084         """Should try to get the contents of a control file using ssh, retry
   1085         get_control_file when getting error.CmdError, raise exception for
   1086         urllib2.HTTPError."""
   1087         dev_server.ImageServerBase.run_call(
   1088                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
   1089         dev_server.ImageServerBase.run_call(
   1090                 mox.IgnoreArg()).AndRaise(E500)
   1091         self.mox.ReplayAll()
   1092         self.assertRaises(dev_server.DevServerException,
   1093                           self.dev_server.get_control_file,
   1094                           '', '')
   1095 
   1096 
   1097     def testGetLatestBuild(self):
   1098         """Should successfully return a build for a given target."""
   1099         self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
   1100         self.mox.StubOutWithMock(dev_server.DevServer, 'devserver_healthy')
   1101 
   1102         dev_server.ImageServer.servers().AndReturn([self._HOST])
   1103         dev_server.ImageServer.devserver_healthy(self._HOST).AndReturn(True)
   1104 
   1105         target = 'x86-generic-release'
   1106         build_string = 'R18-1586.0.0-a1-b1514'
   1107         argument = mox.And(mox.StrContains(self._HOST),
   1108                            mox.StrContains(target))
   1109         dev_server.ImageServerBase.run_call(argument).AndReturn(build_string)
   1110 
   1111         self.mox.ReplayAll()
   1112         build = dev_server.ImageServer.get_latest_build(target)
   1113         self.assertEquals(build_string, build)
   1114 
   1115 
   1116     def testGetLatestBuildWithManyDevservers(self):
   1117         """Should successfully return newest build with multiple devservers."""
   1118         self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
   1119         self.mox.StubOutWithMock(dev_server.DevServer, 'devserver_healthy')
   1120 
   1121         host0_expected = 'http://host0:8080'
   1122         host1_expected = 'http://host1:8082'
   1123 
   1124         dev_server.ImageServer.servers().MultipleTimes().AndReturn(
   1125                 [host0_expected, host1_expected])
   1126 
   1127         dev_server.ImageServer.devserver_healthy(host0_expected).AndReturn(True)
   1128         dev_server.ImageServer.devserver_healthy(host1_expected).AndReturn(True)
   1129 
   1130         target = 'x86-generic-release'
   1131         build_string1 = 'R9-1586.0.0-a1-b1514'
   1132         build_string2 = 'R19-1586.0.0-a1-b3514'
   1133         argument1 = mox.And(mox.StrContains(host0_expected),
   1134                             mox.StrContains(target))
   1135         argument2 = mox.And(mox.StrContains(host1_expected),
   1136                             mox.StrContains(target))
   1137         dev_server.ImageServerBase.run_call(argument1).AndReturn(build_string1)
   1138         dev_server.ImageServerBase.run_call(argument2).AndReturn(build_string2)
   1139 
   1140         self.mox.ReplayAll()
   1141         build = dev_server.ImageServer.get_latest_build(target)
   1142         self.assertEquals(build_string2, build)
   1143 
   1144 
   1145     def testCrashesAreSetToTheCrashServer(self):
   1146         """Should send symbolicate dump rpc calls to crash_server."""
   1147         self.mox.ReplayAll()
   1148         call = self.crash_server.build_call('symbolicate_dump')
   1149         self.assertTrue(call.startswith(self._CRASH_HOST))
   1150 
   1151 
   1152     def _stageTestHelper(self, artifacts=[], files=[], archive_url=None):
   1153         """Helper to test combos of files/artifacts/urls with stage call."""
   1154         expected_archive_url = archive_url
   1155         if not archive_url:
   1156             expected_archive_url = 'gs://my_default_url'
   1157             self.mox.StubOutWithMock(dev_server, '_get_image_storage_server')
   1158             dev_server._get_image_storage_server().AndReturn(
   1159                 'gs://my_default_url')
   1160             name = 'fake/image'
   1161         else:
   1162             # This is embedded in the archive_url. Not needed.
   1163             name = ''
   1164 
   1165         argument1 = mox.And(mox.StrContains(expected_archive_url),
   1166                             mox.StrContains(name),
   1167                             mox.StrContains('artifacts=%s' %
   1168                                             ','.join(artifacts)),
   1169                             mox.StrContains('files=%s' % ','.join(files)),
   1170                             mox.StrContains('stage?'))
   1171         argument2 = mox.And(mox.StrContains(expected_archive_url),
   1172                             mox.StrContains(name),
   1173                             mox.StrContains('artifacts=%s' %
   1174                                             ','.join(artifacts)),
   1175                             mox.StrContains('files=%s' % ','.join(files)),
   1176                             mox.StrContains('is_staged'))
   1177         dev_server.ImageServerBase.run_call(argument1).AndReturn('Success')
   1178         dev_server.ImageServerBase.run_call(argument2).AndReturn('True')
   1179 
   1180         self.mox.ReplayAll()
   1181         self.dev_server.stage_artifacts(name, artifacts, files, archive_url)
   1182         self.mox.VerifyAll()
   1183 
   1184 
   1185     def testStageArtifactsBasic(self):
   1186         """Basic functionality to stage artifacts (similar to
   1187         trigger_download)."""
   1188         self._stageTestHelper(artifacts=['full_payload', 'stateful'])
   1189 
   1190 
   1191     def testStageArtifactsBasicWithFiles(self):
   1192         """Basic functionality to stage artifacts (similar to
   1193         trigger_download)."""
   1194         self._stageTestHelper(artifacts=['full_payload', 'stateful'],
   1195                               files=['taco_bell.coupon'])
   1196 
   1197 
   1198     def testStageArtifactsOnlyFiles(self):
   1199         """Test staging of only file artifacts."""
   1200         self._stageTestHelper(files=['tasty_taco_bell.coupon'])
   1201 
   1202 
   1203     def testStageWithArchiveURL(self):
   1204         """Basic functionality to stage artifacts (similar to
   1205         trigger_download)."""
   1206         self._stageTestHelper(files=['tasty_taco_bell.coupon'],
   1207                               archive_url='gs://tacos_galore/my/dir')
   1208 
   1209 
   1210     def testStagedFileUrl(self):
   1211         """Sanity tests that the staged file url looks right."""
   1212         devserver_label = 'x86-mario-release/R30-1234.0.0'
   1213         url = self.dev_server.get_staged_file_url('stateful.tgz',
   1214                                                   devserver_label)
   1215         expected_url = '/'.join([self._HOST, 'static', devserver_label,
   1216                                  'stateful.tgz'])
   1217         self.assertEquals(url, expected_url)
   1218 
   1219         devserver_label = 'something_crazy/that/you_MIGHT/hate'
   1220         url = self.dev_server.get_staged_file_url('chromiumos_image.bin',
   1221                                                   devserver_label)
   1222         expected_url = '/'.join([self._HOST, 'static', devserver_label,
   1223                                  'chromiumos_image.bin'])
   1224         self.assertEquals(url, expected_url)
   1225 
   1226 
   1227     def _StageTimeoutHelper(self):
   1228         """Helper class for testing staging timeout."""
   1229         self.mox.StubOutWithMock(dev_server.ImageServer, 'call_and_wait')
   1230         dev_server.ImageServer.call_and_wait(
   1231                 call_name='stage',
   1232                 artifacts=mox.IgnoreArg(),
   1233                 files=mox.IgnoreArg(),
   1234                 archive_url=mox.IgnoreArg(),
   1235                 error_message=mox.IgnoreArg()).AndRaise(bin_utils.TimeoutError())
   1236 
   1237 
   1238     def test_StageArtifactsTimeout(self):
   1239         """Test DevServerException is raised when stage_artifacts timed out."""
   1240         self._StageTimeoutHelper()
   1241         self.mox.ReplayAll()
   1242         self.assertRaises(dev_server.DevServerException,
   1243                           self.dev_server.stage_artifacts,
   1244                           image='fake/image', artifacts=['full_payload'])
   1245         self.mox.VerifyAll()
   1246 
   1247 
   1248     def test_TriggerDownloadTimeout(self):
   1249         """Test DevServerException is raised when trigger_download timed out."""
   1250         self._StageTimeoutHelper()
   1251         self.mox.ReplayAll()
   1252         self.assertRaises(dev_server.DevServerException,
   1253                           self.dev_server.trigger_download,
   1254                           image='fake/image')
   1255         self.mox.VerifyAll()
   1256 
   1257 
   1258     def test_FinishDownloadTimeout(self):
   1259         """Test DevServerException is raised when finish_download timed out."""
   1260         self._StageTimeoutHelper()
   1261         self.mox.ReplayAll()
   1262         self.assertRaises(dev_server.DevServerException,
   1263                           self.dev_server.finish_download,
   1264                           image='fake/image')
   1265         self.mox.VerifyAll()
   1266 
   1267 
   1268     def test_compare_load(self):
   1269         """Test load comparison logic.
   1270         """
   1271         load_high_cpu = {'devserver': 'http://devserver_1:8082',
   1272                          dev_server.DevServer.CPU_LOAD: 100.0,
   1273                          dev_server.DevServer.NETWORK_IO: 1024*1024*1.0,
   1274                          dev_server.DevServer.DISK_IO: 1024*1024.0}
   1275         load_high_network = {'devserver': 'http://devserver_1:8082',
   1276                              dev_server.DevServer.CPU_LOAD: 1.0,
   1277                              dev_server.DevServer.NETWORK_IO: 1024*1024*100.0,
   1278                              dev_server.DevServer.DISK_IO: 1024*1024*1.0}
   1279         load_1 = {'devserver': 'http://devserver_1:8082',
   1280                   dev_server.DevServer.CPU_LOAD: 1.0,
   1281                   dev_server.DevServer.NETWORK_IO: 1024*1024*1.0,
   1282                   dev_server.DevServer.DISK_IO: 1024*1024*2.0}
   1283         load_2 = {'devserver': 'http://devserver_1:8082',
   1284                   dev_server.DevServer.CPU_LOAD: 1.0,
   1285                   dev_server.DevServer.NETWORK_IO: 1024*1024*1.0,
   1286                   dev_server.DevServer.DISK_IO: 1024*1024*1.0}
   1287         self.assertFalse(dev_server._is_load_healthy(load_high_cpu))
   1288         self.assertFalse(dev_server._is_load_healthy(load_high_network))
   1289         self.assertTrue(dev_server._compare_load(load_1, load_2) > 0)
   1290 
   1291 
   1292     def _testSuccessfulTriggerDownloadAndroid(self, synchronous=True):
   1293         """Call the dev server's download method with given synchronous
   1294         setting.
   1295 
   1296         @param synchronous: True to call the download method synchronously.
   1297         """
   1298         target = 'test_target'
   1299         branch = 'test_branch'
   1300         build_id = '123456'
   1301         artifacts = android_utils.AndroidArtifacts.get_artifacts_for_reimage(
   1302                 None)
   1303         self.mox.StubOutWithMock(dev_server.AndroidBuildServer,
   1304                                  '_finish_download')
   1305         argument1 = mox.And(mox.StrContains(self._HOST),
   1306                             mox.StrContains(target),
   1307                             mox.StrContains(branch),
   1308                             mox.StrContains(build_id),
   1309                             mox.StrContains('stage?'))
   1310         argument2 = mox.And(mox.StrContains(self._HOST),
   1311                             mox.StrContains(target),
   1312                             mox.StrContains(branch),
   1313                             mox.StrContains(build_id),
   1314                             mox.StrContains('is_staged'))
   1315         dev_server.ImageServerBase.run_call(argument1).AndReturn('Success')
   1316         dev_server.ImageServerBase.run_call(argument2).AndReturn('True')
   1317 
   1318         if synchronous:
   1319             android_build_info = {'target': target,
   1320                                   'build_id': build_id,
   1321                                   'branch': branch}
   1322             build = dev_server.ANDROID_BUILD_NAME_PATTERN % android_build_info
   1323             self.android_dev_server._finish_download(
   1324                     build, artifacts, '', target=target, build_id=build_id,
   1325                     branch=branch)
   1326 
   1327         # Synchronous case requires a call to finish download.
   1328         self.mox.ReplayAll()
   1329         self.android_dev_server.trigger_download(
   1330                 synchronous=synchronous, target=target, build_id=build_id,
   1331                 branch=branch)
   1332         self.mox.VerifyAll()
   1333 
   1334 
   1335     def testSuccessfulTriggerDownloadAndroidSync(self):
   1336         """Call the dev server's download method with synchronous=True."""
   1337         self._testSuccessfulTriggerDownloadAndroid(synchronous=True)
   1338 
   1339 
   1340     def testSuccessfulTriggerDownloadAndroidAsync(self):
   1341         """Call the dev server's download method with synchronous=False."""
   1342         self._testSuccessfulTriggerDownloadAndroid(synchronous=False)
   1343 
   1344 
   1345     def testGetUnrestrictedDevservers(self):
   1346         """Test method get_unrestricted_devservers works as expected."""
   1347         restricted_devserver = 'http://192.168.0.100:8080'
   1348         unrestricted_devserver = 'http://172.1.1.3:8080'
   1349         self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
   1350         dev_server.ImageServer.servers().AndReturn([restricted_devserver,
   1351                                                     unrestricted_devserver])
   1352         self.mox.ReplayAll()
   1353         self.assertEqual(dev_server.ImageServer.get_unrestricted_devservers(
   1354                                 [('192.168.0.0', 24)]),
   1355                          [unrestricted_devserver])
   1356 
   1357 
   1358     def testDevserverHealthy(self):
   1359         """Test which types of connections that method devserver_healthy uses
   1360         for different types of DevServer.
   1361 
   1362         CrashServer always adopts DevServer.run_call.
   1363         ImageServer and AndroidBuildServer use ImageServerBase.run_call.
   1364         """
   1365         argument = mox.StrContains(self._HOST)
   1366 
   1367         # for testing CrashServer
   1368         self.mox.StubOutWithMock(dev_server.DevServer, 'run_call')
   1369         dev_server.DevServer.run_call(
   1370                 argument, timeout=mox.IgnoreArg()).AndReturn(
   1371                         '{"free_disk": 1024}')
   1372         # for testing ImageServer
   1373         dev_server.ImageServerBase.run_call(
   1374                 argument, timeout=mox.IgnoreArg()).AndReturn(
   1375                         '{"free_disk": 1024}')
   1376         # for testing AndroidBuildServer
   1377         dev_server.ImageServerBase.run_call(
   1378                 argument, timeout=mox.IgnoreArg()).AndReturn(
   1379                         '{"free_disk": 1024}')
   1380 
   1381         self.mox.ReplayAll()
   1382         self.assertTrue(dev_server.CrashServer.devserver_healthy(self._HOST))
   1383         self.assertTrue(dev_server.ImageServer.devserver_healthy(self._HOST))
   1384         self.assertTrue(
   1385                 dev_server.AndroidBuildServer.devserver_healthy(self._HOST))
   1386 
   1387 
   1388     def testLocateFile(self):
   1389         """Test locating files for AndriodBuildServer."""
   1390         file_name = 'fake_file'
   1391         artifacts=['full_payload', 'stateful']
   1392         build = 'fake_build'
   1393         argument = mox.And(mox.StrContains(file_name),
   1394                             mox.StrContains(build),
   1395                             mox.StrContains('locate_file'))
   1396         dev_server.ImageServerBase.run_call(argument).AndReturn('file_path')
   1397 
   1398         self.mox.ReplayAll()
   1399         file_location = 'http://nothing/static/fake_build/file_path'
   1400         self.assertEqual(self.android_dev_server.locate_file(
   1401                 file_name, artifacts, build, None), file_location)
   1402 
   1403     def testCmdErrorLocateFile(self):
   1404         """Test locating files for AndriodBuildServer for retry
   1405         error.CmdError, and raise urllib2.URLError."""
   1406         dev_server.ImageServerBase.run_call(
   1407                 mox.IgnoreArg()).AndRaise(CMD_ERROR)
   1408         dev_server.ImageServerBase.run_call(
   1409                 mox.IgnoreArg()).AndRaise(E500)
   1410         self.mox.ReplayAll()
   1411         self.assertRaises(dev_server.DevServerException,
   1412                           self.dev_server.trigger_download,
   1413                           '')
   1414 
   1415 
   1416     def testGetAvailableDevserversForCrashServer(self):
   1417         """Test method get_available_devservers for CrashServer."""
   1418         crash_servers = ['http://crash_servers1:8080']
   1419         host = '127.0.0.1'
   1420         self.mox.StubOutWithMock(dev_server.CrashServer, 'servers')
   1421         dev_server.CrashServer.servers().AndReturn(crash_servers)
   1422         self.mox.ReplayAll()
   1423         self.assertEqual(dev_server.CrashServer.get_available_devservers(host),
   1424                         (crash_servers, False))
   1425 
   1426 
   1427     def testGetAvailableDevserversForImageServer(self):
   1428         """Test method get_available_devservers for ImageServer."""
   1429         unrestricted_host = '100.0.0.99'
   1430         unrestricted_servers = ['http://100.0.0.10:8080',
   1431                                 'http://128.0.0.10:8080']
   1432         same_subnet_unrestricted_servers = ['http://100.0.0.10:8080']
   1433         restricted_host = '127.0.0.99'
   1434         restricted_servers = ['http://127.0.0.10:8080']
   1435         all_servers = unrestricted_servers + restricted_servers
   1436         # Set restricted subnets
   1437         restricted_subnets = [('127.0.0.0', 24)]
   1438         self.mox.StubOutWithMock(dev_server.ImageServerBase, 'servers')
   1439         dev_server.ImageServerBase.servers().MultipleTimes().AndReturn(
   1440                 all_servers)
   1441         self.mox.ReplayAll()
   1442         # dut in unrestricted subnet shall be offered devserver in the same
   1443         # subnet first, and allow retry.
   1444         self.assertEqual(
   1445                 dev_server.ImageServer.get_available_devservers(
   1446                         unrestricted_host, True, restricted_subnets),
   1447                 (same_subnet_unrestricted_servers, True))
   1448 
   1449         # If prefer_local_devserver is set to False, allow any devserver in
   1450         # unrestricted subet to be available, and retry is not allowed.
   1451         self.assertEqual(
   1452                 dev_server.ImageServer.get_available_devservers(
   1453                         unrestricted_host, False, restricted_subnets),
   1454                 (unrestricted_servers, False))
   1455 
   1456         # When no hostname is specified, all devservers in unrestricted subnets
   1457         # should be considered, and retry is not allowed.
   1458         self.assertEqual(
   1459                 dev_server.ImageServer.get_available_devservers(
   1460                         None, True, restricted_subnets),
   1461                 (unrestricted_servers, False))
   1462 
   1463         # dut in restricted subnet should only be offered devserver in the
   1464         # same restricted subnet, and retry is not allowed.
   1465         self.assertEqual(
   1466                 dev_server.ImageServer.get_available_devservers(
   1467                         restricted_host, True, restricted_subnets),
   1468                 (restricted_servers, False))
   1469 
   1470 
   1471 if __name__ == "__main__":
   1472     unittest.main()
   1473