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
9import httplib
10import mox
11import StringIO
12import time
13import unittest
14import urllib2
15
16from autotest_lib.client.common_lib import error
17from autotest_lib.client.common_lib import global_config
18from autotest_lib.client.common_lib.cros import dev_server
19from autotest_lib.client.common_lib.cros import retry
20
21def retry_mock(ExceptionToCheck, timeout_min):
22    """A mock retry decorator to use in place of the actual one for testing.
23
24    @param ExceptionToCheck: the exception to check.
25    @param timeout_mins: Amount of time in mins to wait before timing out.
26
27    """
28    def inner_retry(func):
29        """The actual decorator.
30
31        @param func: Function to be called in decorator.
32
33        """
34        return func
35
36    return inner_retry
37
38
39class DevServerTest(mox.MoxTestBase):
40    """Unit tests for dev_server.DevServer.
41
42    @var _HOST: fake dev server host address.
43    """
44
45    _HOST = 'http://nothing'
46    _CRASH_HOST = 'http://nothing-crashed'
47    _CONFIG = global_config.global_config
48
49
50    def setUp(self):
51        super(DevServerTest, self).setUp()
52        self.crash_server = dev_server.CrashServer(DevServerTest._CRASH_HOST)
53        self.dev_server = dev_server.ImageServer(DevServerTest._HOST)
54        self.android_dev_server = dev_server.AndroidBuildServer(
55                DevServerTest._HOST)
56        self.mox.StubOutWithMock(urllib2, 'urlopen')
57        # Hide local restricted_subnets setting.
58        dev_server.RESTRICTED_SUBNETS = []
59
60
61    def testSimpleResolve(self):
62        """One devserver, verify we resolve to it."""
63        self.mox.StubOutWithMock(dev_server, '_get_dev_server_list')
64        self.mox.StubOutWithMock(dev_server.DevServer, 'devserver_healthy')
65        dev_server._get_dev_server_list().MultipleTimes().AndReturn(
66                [DevServerTest._HOST])
67        dev_server.DevServer.devserver_healthy(DevServerTest._HOST).AndReturn(
68                                                                        True)
69        self.mox.ReplayAll()
70        devserver = dev_server.ImageServer.resolve('my_build')
71        self.assertEquals(devserver.url(), DevServerTest._HOST)
72
73
74    def testResolveWithFailure(self):
75        """Ensure we rehash on a failed ping on a bad_host."""
76        self.mox.StubOutWithMock(dev_server, '_get_dev_server_list')
77        bad_host, good_host = 'http://bad_host:99', 'http://good_host:8080'
78        dev_server._get_dev_server_list().MultipleTimes().AndReturn(
79                [bad_host, good_host])
80
81        # Mock out bad ping failure to bad_host by raising devserver exception.
82        urllib2.urlopen(mox.StrContains(bad_host), data=None).AndRaise(
83                dev_server.DevServerException())
84        # Good host is good.
85        to_return = StringIO.StringIO('{"free_disk": 1024}')
86        urllib2.urlopen(mox.StrContains(good_host),
87                        data=None).AndReturn(to_return)
88
89        self.mox.ReplayAll()
90        host = dev_server.ImageServer.resolve(0) # Using 0 as it'll hash to 0.
91        self.assertEquals(host.url(), good_host)
92        self.mox.VerifyAll()
93
94
95    def testResolveWithFailureURLError(self):
96        """Ensure we rehash on a failed ping on a bad_host after urlerror."""
97        # Retry mock just return the original method.
98        retry.retry = retry_mock
99        self.mox.StubOutWithMock(dev_server, '_get_dev_server_list')
100        bad_host, good_host = 'http://bad_host:99', 'http://good_host:8080'
101        dev_server._get_dev_server_list().MultipleTimes().AndReturn(
102                [bad_host, good_host])
103
104        # Mock out bad ping failure to bad_host by raising devserver exception.
105        urllib2.urlopen(mox.StrContains(bad_host),
106                data=None).MultipleTimes().AndRaise(
107                        urllib2.URLError('urlopen connection timeout'))
108
109        # Good host is good.
110        to_return = StringIO.StringIO('{"free_disk": 1024}')
111        urllib2.urlopen(mox.StrContains(good_host),
112                data=None).AndReturn(to_return)
113
114        self.mox.ReplayAll()
115        host = dev_server.ImageServer.resolve(0) # Using 0 as it'll hash to 0.
116        self.assertEquals(host.url(), good_host)
117        self.mox.VerifyAll()
118
119
120    def testResolveWithManyDevservers(self):
121        """Should be able to return different urls with multiple devservers."""
122        self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
123        self.mox.StubOutWithMock(dev_server.DevServer, 'devserver_healthy')
124
125        host0_expected = 'http://host0:8080'
126        host1_expected = 'http://host1:8082'
127
128        dev_server.ImageServer.servers().MultipleTimes().AndReturn(
129                [host0_expected, host1_expected])
130        dev_server.DevServer.devserver_healthy(host0_expected).AndReturn(True)
131        dev_server.DevServer.devserver_healthy(host1_expected).AndReturn(True)
132
133        self.mox.ReplayAll()
134        host0 = dev_server.ImageServer.resolve(0)
135        host1 = dev_server.ImageServer.resolve(1)
136        self.mox.VerifyAll()
137
138        self.assertEqual(host0.url(), host0_expected)
139        self.assertEqual(host1.url(), host1_expected)
140
141
142    def _returnHttpServerError(self):
143        e500 = urllib2.HTTPError(url='',
144                                 code=httplib.INTERNAL_SERVER_ERROR,
145                                 msg='',
146                                 hdrs=None,
147                                 fp=StringIO.StringIO('Expected.'))
148        urllib2.urlopen(mox.IgnoreArg()).AndRaise(e500)
149
150
151    def _returnHttpForbidden(self):
152        e403 = urllib2.HTTPError(url='',
153                                 code=httplib.FORBIDDEN,
154                                 msg='',
155                                 hdrs=None,
156                                 fp=StringIO.StringIO('Expected.'))
157        urllib2.urlopen(mox.IgnoreArg()).AndRaise(e403)
158
159
160    def testSuccessfulTriggerDownloadSync(self):
161        """Call the dev server's download method with synchronous=True."""
162        name = 'fake/image'
163        self.mox.StubOutWithMock(dev_server.ImageServer, '_finish_download')
164        to_return = StringIO.StringIO('Success')
165        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
166                                mox.StrContains(name),
167                                mox.StrContains('stage?'))).AndReturn(to_return)
168        to_return = StringIO.StringIO('True')
169        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
170                                mox.StrContains(name),
171                                mox.StrContains('is_staged'))).AndReturn(
172                                                                      to_return)
173        self.dev_server._finish_download(name, mox.IgnoreArg(), mox.IgnoreArg())
174
175        # Synchronous case requires a call to finish download.
176        self.mox.ReplayAll()
177        self.dev_server.trigger_download(name, synchronous=True)
178        self.mox.VerifyAll()
179
180
181    def testSuccessfulTriggerDownloadASync(self):
182        """Call the dev server's download method with synchronous=False."""
183        name = 'fake/image'
184        to_return = StringIO.StringIO('Success')
185        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
186                                mox.StrContains(name),
187                                mox.StrContains('stage?'))).AndReturn(to_return)
188        to_return = StringIO.StringIO('True')
189        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
190                                mox.StrContains(name),
191                                mox.StrContains('is_staged'))).AndReturn(
192                                                                      to_return)
193
194        self.mox.ReplayAll()
195        self.dev_server.trigger_download(name, synchronous=False)
196        self.mox.VerifyAll()
197
198
199    def testURLErrorRetryTriggerDownload(self):
200        """Should retry on URLError, but pass through real exception."""
201        self.mox.StubOutWithMock(time, 'sleep')
202
203        refused = urllib2.URLError('[Errno 111] Connection refused')
204        urllib2.urlopen(mox.IgnoreArg()).AndRaise(refused)
205        time.sleep(mox.IgnoreArg())
206        self._returnHttpForbidden()
207        self.mox.ReplayAll()
208        self.assertRaises(dev_server.DevServerException,
209                          self.dev_server.trigger_download,
210                          '')
211
212
213    def testErrorTriggerDownload(self):
214        """Should call the dev server's download method, fail gracefully."""
215        self._returnHttpServerError()
216        self.mox.ReplayAll()
217        self.assertRaises(dev_server.DevServerException,
218                          self.dev_server.trigger_download,
219                          '')
220
221
222    def testForbiddenTriggerDownload(self):
223        """Should call the dev server's download method, get exception."""
224        self._returnHttpForbidden()
225        self.mox.ReplayAll()
226        self.assertRaises(dev_server.DevServerException,
227                          self.dev_server.trigger_download,
228                          '')
229
230
231    def testSuccessfulFinishDownload(self):
232        """Should successfully call the dev server's finish download method."""
233        name = 'fake/image'
234        to_return = StringIO.StringIO('Success')
235        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
236                                mox.StrContains(name),
237                                mox.StrContains('stage?'))).AndReturn(to_return)
238        to_return = StringIO.StringIO('True')
239        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
240                                mox.StrContains(name),
241                                mox.StrContains('is_staged'))).AndReturn(
242                                                                      to_return)
243
244        # Synchronous case requires a call to finish download.
245        self.mox.ReplayAll()
246        self.dev_server.finish_download(name)  # Raises on failure.
247        self.mox.VerifyAll()
248
249
250    def testErrorFinishDownload(self):
251        """Should call the dev server's finish download method, fail gracefully.
252        """
253        self._returnHttpServerError()
254        self.mox.ReplayAll()
255        self.assertRaises(dev_server.DevServerException,
256                          self.dev_server.finish_download,
257                          '')
258
259
260    def testListControlFiles(self):
261        """Should successfully list control files from the dev server."""
262        name = 'fake/build'
263        control_files = ['file/one', 'file/two']
264        to_return = StringIO.StringIO('\n'.join(control_files))
265        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
266                                mox.StrContains(name))).AndReturn(to_return)
267        self.mox.ReplayAll()
268        paths = self.dev_server.list_control_files(name)
269        self.assertEquals(len(paths), 2)
270        for f in control_files:
271            self.assertTrue(f in paths)
272
273
274    def testFailedListControlFiles(self):
275        """Should call the dev server's list-files method, get exception."""
276        self._returnHttpServerError()
277        self.mox.ReplayAll()
278        self.assertRaises(dev_server.DevServerException,
279                          self.dev_server.list_control_files,
280                          '')
281
282
283    def testExplodingListControlFiles(self):
284        """Should call the dev server's list-files method, get exception."""
285        self._returnHttpForbidden()
286        self.mox.ReplayAll()
287        self.assertRaises(dev_server.DevServerException,
288                          self.dev_server.list_control_files,
289                          '')
290
291
292    def testGetControlFile(self):
293        """Should successfully get a control file from the dev server."""
294        name = 'fake/build'
295        file = 'file/one'
296        contents = 'Multi-line\nControl File Contents\n'
297        to_return = StringIO.StringIO(contents)
298        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
299                                mox.StrContains(name),
300                                mox.StrContains(file))).AndReturn(to_return)
301        self.mox.ReplayAll()
302        self.assertEquals(self.dev_server.get_control_file(name, file),
303                          contents)
304
305
306    def testErrorGetControlFile(self):
307        """Should try to get the contents of a control file, get exception."""
308        self._returnHttpServerError()
309        self.mox.ReplayAll()
310        self.assertRaises(dev_server.DevServerException,
311                          self.dev_server.get_control_file,
312                          '', '')
313
314
315    def testForbiddenGetControlFile(self):
316        """Should try to get the contents of a control file, get exception."""
317        self._returnHttpForbidden()
318        self.mox.ReplayAll()
319        self.assertRaises(dev_server.DevServerException,
320                          self.dev_server.get_control_file,
321                          '', '')
322
323
324    def testGetLatestBuild(self):
325        """Should successfully return a build for a given target."""
326        self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
327        self.mox.StubOutWithMock(dev_server.DevServer, 'devserver_healthy')
328
329        dev_server.ImageServer.servers().AndReturn([self._HOST])
330        dev_server.DevServer.devserver_healthy(self._HOST).AndReturn(True)
331
332        target = 'x86-generic-release'
333        build_string = 'R18-1586.0.0-a1-b1514'
334        to_return = StringIO.StringIO(build_string)
335        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
336                                mox.StrContains(target))).AndReturn(to_return)
337        self.mox.ReplayAll()
338        build = dev_server.ImageServer.get_latest_build(target)
339        self.assertEquals(build_string, build)
340
341
342    def testGetLatestBuildWithManyDevservers(self):
343        """Should successfully return newest build with multiple devservers."""
344        self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
345        self.mox.StubOutWithMock(dev_server.DevServer, 'devserver_healthy')
346
347        host0_expected = 'http://host0:8080'
348        host1_expected = 'http://host1:8082'
349
350        dev_server.ImageServer.servers().MultipleTimes().AndReturn(
351                [host0_expected, host1_expected])
352
353        dev_server.DevServer.devserver_healthy(host0_expected).AndReturn(True)
354        dev_server.DevServer.devserver_healthy(host1_expected).AndReturn(True)
355
356        target = 'x86-generic-release'
357        build_string1 = 'R9-1586.0.0-a1-b1514'
358        build_string2 = 'R19-1586.0.0-a1-b3514'
359        to_return1 = StringIO.StringIO(build_string1)
360        to_return2 = StringIO.StringIO(build_string2)
361        urllib2.urlopen(mox.And(mox.StrContains(host0_expected),
362                                mox.StrContains(target))).AndReturn(to_return1)
363        urllib2.urlopen(mox.And(mox.StrContains(host1_expected),
364                                mox.StrContains(target))).AndReturn(to_return2)
365
366        self.mox.ReplayAll()
367        build = dev_server.ImageServer.get_latest_build(target)
368        self.assertEquals(build_string2, build)
369
370
371    def testCrashesAreSetToTheCrashServer(self):
372        """Should send symbolicate dump rpc calls to crash_server."""
373        self.mox.ReplayAll()
374        call = self.crash_server.build_call('symbolicate_dump')
375        self.assertTrue(call.startswith(self._CRASH_HOST))
376
377
378    def _stageTestHelper(self, artifacts=[], files=[], archive_url=None):
379        """Helper to test combos of files/artifacts/urls with stage call."""
380        expected_archive_url = archive_url
381        if not archive_url:
382            expected_archive_url = 'gs://my_default_url'
383            self.mox.StubOutWithMock(dev_server, '_get_image_storage_server')
384            dev_server._get_image_storage_server().AndReturn(
385                'gs://my_default_url')
386            name = 'fake/image'
387        else:
388            # This is embedded in the archive_url. Not needed.
389            name = ''
390
391        to_return = StringIO.StringIO('Success')
392        urllib2.urlopen(mox.And(mox.StrContains(expected_archive_url),
393                                mox.StrContains(name),
394                                mox.StrContains('artifacts=%s' %
395                                                ','.join(artifacts)),
396                                mox.StrContains('files=%s' % ','.join(files)),
397                                mox.StrContains('stage?'))).AndReturn(to_return)
398        to_return = StringIO.StringIO('True')
399        urllib2.urlopen(mox.And(mox.StrContains(expected_archive_url),
400                                mox.StrContains(name),
401                                mox.StrContains('artifacts=%s' %
402                                                ','.join(artifacts)),
403                                mox.StrContains('files=%s' % ','.join(files)),
404                                mox.StrContains('is_staged'))).AndReturn(
405                                        to_return)
406
407        self.mox.ReplayAll()
408        self.dev_server.stage_artifacts(name, artifacts, files, archive_url)
409        self.mox.VerifyAll()
410
411
412    def testStageArtifactsBasic(self):
413        """Basic functionality to stage artifacts (similar to trigger_download).
414        """
415        self._stageTestHelper(artifacts=['full_payload', 'stateful'])
416
417
418    def testStageArtifactsBasicWithFiles(self):
419        """Basic functionality to stage artifacts (similar to trigger_download).
420        """
421        self._stageTestHelper(artifacts=['full_payload', 'stateful'],
422                              files=['taco_bell.coupon'])
423
424
425    def testStageArtifactsOnlyFiles(self):
426        """Test staging of only file artifacts."""
427        self._stageTestHelper(files=['tasty_taco_bell.coupon'])
428
429
430    def testStageWithArchiveURL(self):
431        """Basic functionality to stage artifacts (similar to trigger_download).
432        """
433        self._stageTestHelper(files=['tasty_taco_bell.coupon'],
434                              archive_url='gs://tacos_galore/my/dir')
435
436
437    def testStagedFileUrl(self):
438        """Sanity tests that the staged file url looks right."""
439        devserver_label = 'x86-mario-release/R30-1234.0.0'
440        url = self.dev_server.get_staged_file_url('stateful.tgz',
441                                                  devserver_label)
442        expected_url = '/'.join([self._HOST, 'static', devserver_label,
443                                 'stateful.tgz'])
444        self.assertEquals(url, expected_url)
445
446        devserver_label = 'something_crazy/that/you_MIGHT/hate'
447        url = self.dev_server.get_staged_file_url('chromiumos_image.bin',
448                                                  devserver_label)
449        expected_url = '/'.join([self._HOST, 'static', devserver_label,
450                                 'chromiumos_image.bin'])
451        self.assertEquals(url, expected_url)
452
453
454    def _StageTimeoutHelper(self):
455        """Helper class for testing staging timeout."""
456        self.mox.StubOutWithMock(dev_server.ImageServer, 'call_and_wait')
457        dev_server.ImageServer.call_and_wait(
458                call_name='stage',
459                artifacts=mox.IgnoreArg(),
460                files=mox.IgnoreArg(),
461                archive_url=mox.IgnoreArg(),
462                error_message=mox.IgnoreArg()).AndRaise(error.TimeoutException)
463
464
465    def test_StageArtifactsTimeout(self):
466        """Test DevServerException is raised when stage_artifacts timed out."""
467        self._StageTimeoutHelper()
468        self.mox.ReplayAll()
469        self.assertRaises(dev_server.DevServerException,
470                          self.dev_server.stage_artifacts,
471                          image='fake/image', artifacts=['full_payload'])
472        self.mox.VerifyAll()
473
474
475    def test_TriggerDownloadTimeout(self):
476        """Test DevServerException is raised when trigger_download timed out."""
477        self._StageTimeoutHelper()
478        self.mox.ReplayAll()
479        self.assertRaises(dev_server.DevServerException,
480                          self.dev_server.trigger_download,
481                          image='fake/image')
482        self.mox.VerifyAll()
483
484
485    def test_FinishDownloadTimeout(self):
486        """Test DevServerException is raised when finish_download timed out."""
487        self._StageTimeoutHelper()
488        self.mox.ReplayAll()
489        self.assertRaises(dev_server.DevServerException,
490                          self.dev_server.finish_download,
491                          image='fake/image')
492        self.mox.VerifyAll()
493
494
495    def test_compare_load(self):
496        """Test load comparison logic.
497        """
498        load_high_cpu = {'devserver': 'http://devserver_1:8082',
499                         dev_server.DevServer.CPU_LOAD: 100.0,
500                         dev_server.DevServer.NETWORK_IO: 1024*1024*1.0,
501                         dev_server.DevServer.DISK_IO: 1024*1024.0}
502        load_high_network = {'devserver': 'http://devserver_1:8082',
503                             dev_server.DevServer.CPU_LOAD: 1.0,
504                             dev_server.DevServer.NETWORK_IO: 1024*1024*100.0,
505                             dev_server.DevServer.DISK_IO: 1024*1024*1.0}
506        load_1 = {'devserver': 'http://devserver_1:8082',
507                  dev_server.DevServer.CPU_LOAD: 1.0,
508                  dev_server.DevServer.NETWORK_IO: 1024*1024*1.0,
509                  dev_server.DevServer.DISK_IO: 1024*1024*2.0}
510        load_2 = {'devserver': 'http://devserver_1:8082',
511                  dev_server.DevServer.CPU_LOAD: 1.0,
512                  dev_server.DevServer.NETWORK_IO: 1024*1024*1.0,
513                  dev_server.DevServer.DISK_IO: 1024*1024*1.0}
514        self.assertFalse(dev_server._is_load_healthy(load_high_cpu))
515        self.assertFalse(dev_server._is_load_healthy(load_high_network))
516        self.assertTrue(dev_server._compare_load(load_1, load_2) > 0)
517
518
519    def _testSuccessfulTriggerDownloadAndroid(self, synchronous=True):
520        """Call the dev server's download method with given synchronous setting.
521
522        @param synchronous: True to call the download method synchronously.
523        """
524        target = 'test_target'
525        branch = 'test_branch'
526        build_id = '123456'
527        self.mox.StubOutWithMock(dev_server.AndroidBuildServer,
528                                 '_finish_download')
529        to_return = StringIO.StringIO('Success')
530        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
531                                mox.StrContains(target),
532                                mox.StrContains(branch),
533                                mox.StrContains(build_id),
534                                mox.StrContains('stage?'))).AndReturn(to_return)
535        to_return = StringIO.StringIO('True')
536        urllib2.urlopen(mox.And(mox.StrContains(self._HOST),
537                                mox.StrContains(target),
538                                mox.StrContains(branch),
539                                mox.StrContains(build_id),
540                                mox.StrContains('is_staged'))).AndReturn(
541                                                                      to_return)
542        if synchronous:
543            android_build_info = {'target': target,
544                                  'build_id': build_id,
545                                  'branch': branch}
546            build = dev_server.ANDROID_BUILD_NAME_PATTERN % android_build_info
547            self.android_dev_server._finish_download(
548                    build,
549                    dev_server._ANDROID_ARTIFACTS_TO_BE_STAGED_FOR_IMAGE, '',
550                    target=target, build_id=build_id, branch=branch)
551
552        # Synchronous case requires a call to finish download.
553        self.mox.ReplayAll()
554        self.android_dev_server.trigger_download(
555                synchronous=synchronous, target=target, build_id=build_id,
556                branch=branch)
557        self.mox.VerifyAll()
558
559
560    def testSuccessfulTriggerDownloadAndroidSync(self):
561        """Call the dev server's download method with synchronous=True."""
562        self._testSuccessfulTriggerDownloadAndroid(synchronous=True)
563
564
565    def testSuccessfulTriggerDownloadAndroidAsync(self):
566        """Call the dev server's download method with synchronous=False."""
567        self._testSuccessfulTriggerDownloadAndroid(synchronous=False)
568
569
570    def testGetUnrestrictedDevservers(self):
571        """Test method get_unrestricted_devservers works as expected."""
572        restricted_devserver = 'http://192.168.0.100:8080'
573        unrestricted_devserver = 'http://172.1.1.3:8080'
574        self.mox.StubOutWithMock(dev_server.ImageServer, 'servers')
575        dev_server.ImageServer.servers().AndReturn([restricted_devserver,
576                                                    unrestricted_devserver])
577        self.mox.ReplayAll()
578        self.assertEqual(dev_server.ImageServer.get_unrestricted_devservers(
579                                [('192.168.0.0', 24)]),
580                         [unrestricted_devserver])
581
582
583if __name__ == "__main__":
584    unittest.main()
585