src/tools/swarming_client/tests/swarming_test.py

   1 #!/usr/bin/env python
   2 # Copyright 2013 The Swarming Authors. All rights reserved.
   3 # Use of this source code is governed under the Apache License, Version 2.0 that
   4 # can be found in the LICENSE file.
   5
   6 import datetime
   7 import getpass
   8 import hashlib
   9 import json
  10 import logging
  11 import os
  12 import shutil
  13 import StringIO
  14 import sys
  15 import tempfile
  16 import threading
  17 import unittest
  18
  19 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  20 sys.path.insert(0, ROOT_DIR)
  21 sys.path.insert(0, os.path.join(ROOT_DIR, 'third_party'))
  22
  23 from depot_tools import auto_stub
  24 import swarming
  25 import test_utils
  26 from utils import net
  27
  28
  29 ALGO = hashlib.sha1
  30 FILE_NAME = u'test.isolated'
  31 FILE_HASH = u'1' * 40
  32 TEST_NAME = u'unit_tests'
  33
  34
  35 TEST_CASE_SUCCESS = (
  36   '[----------] 2 tests from StaticCookiePolicyTest\n'
  37   '[ RUN      ] StaticCookiePolicyTest.AllowAllCookiesTest\n'
  38   '[       OK ] StaticCookiePolicyTest.AllowAllCookiesTest (0 ms)\n'
  39   '[ RUN      ] StaticCookiePolicyTest.BlockAllCookiesTest\n'
  40   '[       OK ] StaticCookiePolicyTest.BlockAllCookiesTest (0 ms)\n'
  41   '[----------] 2 tests from StaticCookiePolicyTest (0 ms total)\n'
  42   '\n'
  43   '[----------] 1 test from TCPListenSocketTest\n'
  44   '[ RUN      ] TCPListenSocketTest.ServerSend\n'
  45   '[       OK ] TCPListenSocketTest.ServerSend (1 ms)\n'
  46   '[----------] 1 test from TCPListenSocketTest (1 ms total)\n')
  47
  48
  49 TEST_CASE_FAILURE = (
  50   '[----------] 2 tests from StaticCookiePolicyTest\n'
  51   '[ RUN      ] StaticCookiePolicyTest.AllowAllCookiesTest\n'
  52   '[       OK ] StaticCookiePolicyTest.AllowAllCookiesTest (0 ms)\n'
  53   '[ RUN      ] StaticCookiePolicyTest.BlockAllCookiesTest\n'
  54   'C:\\win\\build\\src\\chrome\\test.cc: error: Value of: result()\n'
  55   '  Actual: false\n'
  56   'Expected: true\n'
  57   '[  FAILED  ] StaticCookiePolicyTest.BlockAllCookiesTest (0 ms)\n'
  58   '[----------] 2 tests from StaticCookiePolicyTest (0 ms total)\n'
  59   '\n'
  60   '[----------] 1 test from TCPListenSocketTest\n'
  61   '[ RUN      ] TCPListenSocketTest.ServerSend\n'
  62   '[       OK ] TCPListenSocketTest.ServerSend (1 ms)\n'
  63   '[----------] 1 test from TCPListenSocketTest (1 ms total)\n')
  64
  65
  66 SWARM_OUTPUT_SUCCESS = (
  67   '[ RUN      ] unittests.Run Test\n' +
  68   TEST_CASE_SUCCESS +
  69   '[       OK ] unittests.Run Test (2549 ms)\n'
  70   '[ RUN      ] unittests.Clean Up\n'
  71   'No output!\n'
  72   '[       OK ] unittests.Clean Up (6 ms)\n'
  73   '\n'
  74   '[----------] unittests summary\n'
  75   '[==========] 2 tests ran. (2556 ms total)\n')
  76
  77
  78 SWARM_OUTPUT_FAILURE = (
  79   '[ RUN      ] unittests.Run Test\n' +
  80   TEST_CASE_FAILURE +
  81   '[       OK ] unittests.Run Test (2549 ms)\n'
  82   '[ RUN      ] unittests.Clean Up\n'
  83   'No output!\n'
  84   '[       OK ] unittests.Clean Up (6 ms)\n'
  85   '\n'
  86   '[----------] unittests summary\n'
  87   '[==========] 2 tests ran. (2556 ms total)\n')
  88
  89
  90 SWARM_OUTPUT_WITH_NO_TEST_OUTPUT = (
  91   '\n'
  92   'Unable to connection to swarm machine.\n')
  93
  94
  95 TEST_SHARD_1 = 'Note: This is test shard 1 of 3.'
  96 TEST_SHARD_2 = 'Note: This is test shard 2 of 3.'
  97 TEST_SHARD_3 = 'Note: This is test shard 3 of 3.'
  98
  99
 100 SWARM_SHARD_OUTPUT = (
 101   '[ RUN      ] unittests.Run Test\n'
 102   '%s\n'
 103   '[       OK ] unittests.Run Test (2549 ms)\n'
 104   '[ RUN      ] unittests.Clean Up\n'
 105   'No output!\n'
 106   '[       OK ] unittests.Clean Up (6 ms)\n'
 107   '\n'
 108   '[----------] unittests summary\n'
 109   '[==========] 2 tests ran. (2556 ms total)\n')
 110
 111
 112 TEST_SHARD_OUTPUT_1 = SWARM_SHARD_OUTPUT % TEST_SHARD_1
 113 TEST_SHARD_OUTPUT_2 = SWARM_SHARD_OUTPUT % TEST_SHARD_2
 114 TEST_SHARD_OUTPUT_3 = SWARM_SHARD_OUTPUT % TEST_SHARD_3
 115
 116
 117 def gen_data(index, shard_output, exit_codes):
 118   return {
 119     u'config_instance_index': index,
 120     u'exit_codes': unicode(exit_codes),
 121     u'machine_id': u'host',
 122     u'machine_tag': u'localhost',
 123     u'output': unicode(shard_output),
 124   }
 125
 126
 127 def gen_yielded_data(index, shard_output, exit_codes):
 128   """Returns an entry as it would be yielded by yield_results()."""
 129   return index, gen_data(index, shard_output, exit_codes)
 130
 131
 132 def generate_url_response(index, shard_output, exit_codes):
 133   return net.HttpResponse.get_fake_response(
 134       json.dumps(gen_data(index, shard_output, exit_codes)), 'mocked_url')
 135
 136
 137 def get_swarm_results(keys):
 138   """Simplifies the call to yield_results().
 139
 140   The timeout is hard-coded to 10 seconds.
 141   """
 142   return list(swarming.yield_results('http://host:9001', keys, 10., None))
 143
 144
 145 def main(args):
 146   """Bypassies swarming.main()'s exception handling.
 147
 148   It gets in the way when debugging test failures.
 149   """
 150   dispatcher = swarming.subcommand.CommandDispatcher('swarming')
 151   return dispatcher.execute(swarming.OptionParserSwarming(), args)
 152
 153
 154 class TestCase(auto_stub.TestCase):
 155   """Base class that defines the url_open mock."""
 156   def setUp(self):
 157     super(TestCase, self).setUp()
 158     self._lock = threading.Lock()
 159     self.requests = []
 160     self.mock(swarming.net.HttpService, 'request', self._url_open)
 161     self.mock(swarming.time, 'sleep', lambda x: None)
 162     self.mock(swarming.subprocess, 'call', lambda *_: self.fail())
 163     self.mock(sys, 'stdout', StringIO.StringIO())
 164     self.mock(sys, 'stderr', StringIO.StringIO())
 165
 166   def tearDown(self):
 167     try:
 168       if not self.has_failed():
 169         self._check_output('', '')
 170         self.assertEqual([], self.requests)
 171     finally:
 172       super(TestCase, self).tearDown()
 173
 174   def _check_output(self, out, err):
 175     self.assertEqual(out, sys.stdout.getvalue())
 176     self.assertEqual(err, sys.stderr.getvalue())
 177
 178     # Flush their content by mocking them again.
 179     self.mock(sys, 'stdout', StringIO.StringIO())
 180     self.mock(sys, 'stderr', StringIO.StringIO())
 181
 182   def _url_open(self, url, **kwargs):
 183     logging.info('url_open(%s)', url)
 184     # Ignore 'stream' argument, it's not important for these tests.
 185     kwargs.pop('stream')
 186     with self._lock:
 187       # Since the client is multi-threaded, requests can be processed out of
 188       # order.
 189       for index, r in enumerate(self.requests):
 190         if r[0] == url and r[1] == kwargs:
 191           _, _, returned = self.requests.pop(index)
 192           break
 193       else:
 194         self.fail(
 195             'Failed to find url %s\n%s\nRemaining:\n%s' % (
 196               url,
 197               json.dumps(kwargs, indent=2, sort_keys=True),
 198               json.dumps(
 199                   [(i[0], i[1]) for i in self.requests],
 200                   indent=2, sort_keys=True)))
 201     return returned
 202
 203
 204 class TestGetTestKeys(TestCase):
 205   def test_no_keys(self):
 206     self.mock(swarming.time, 'sleep', lambda x: x)
 207     self.requests = [
 208       (
 209         '/get_matching_test_cases?name=my_test',
 210         {'retry_404': True},
 211         StringIO.StringIO('No matching Test Cases'),
 212       ) for _ in range(net.URL_OPEN_MAX_ATTEMPTS)
 213     ]
 214     try:
 215       swarming.get_task_keys('http://host:9001', 'my_test')
 216       self.fail()
 217     except swarming.Failure as e:
 218       msg = (
 219           'Error: Unable to find any task with the name, my_test, on swarming '
 220           'server')
 221       self.assertEqual(msg, e.args[0])
 222
 223   def test_no_keys_on_first_attempt(self):
 224     self.mock(swarming.time, 'sleep', lambda x: x)
 225     keys = ['key_1', 'key_2']
 226     self.requests = [
 227       (
 228         '/get_matching_test_cases?name=my_test',
 229         {'retry_404': True},
 230         StringIO.StringIO('No matching Test Cases'),
 231       ),
 232       (
 233         '/get_matching_test_cases?name=my_test',
 234         {'retry_404': True},
 235         StringIO.StringIO(json.dumps(keys)),
 236       ),
 237     ]
 238     actual = swarming.get_task_keys('http://host:9001', 'my_test')
 239     self.assertEqual(keys, actual)
 240
 241   def test_find_keys(self):
 242     keys = ['key_1', 'key_2']
 243     self.requests = [
 244       (
 245         '/get_matching_test_cases?name=my_test',
 246         {'retry_404': True},
 247         StringIO.StringIO(json.dumps(keys)),
 248       ),
 249     ]
 250     actual = swarming.get_task_keys('http://host:9001', 'my_test')
 251     self.assertEqual(keys, actual)
 252
 253
 254 class TestGetSwarmResults(TestCase):
 255   def test_success(self):
 256     self.requests = [
 257       (
 258         '/get_result?r=key1',
 259         {'retry_404': False, 'retry_50x': False},
 260         generate_url_response(0, SWARM_OUTPUT_SUCCESS, '0, 0'),
 261       ),
 262     ]
 263     expected = [gen_yielded_data(0, SWARM_OUTPUT_SUCCESS, '0, 0')]
 264     actual = get_swarm_results(['key1'])
 265     self.assertEqual(expected, actual)
 266
 267   def test_failure(self):
 268     self.requests = [
 269       (
 270         '/get_result?r=key1',
 271         {'retry_404': False, 'retry_50x': False},
 272         generate_url_response(0, SWARM_OUTPUT_FAILURE, '0, 1'),
 273       ),
 274     ]
 275     expected = [gen_yielded_data(0, SWARM_OUTPUT_FAILURE, '0, 1')]
 276     actual = get_swarm_results(['key1'])
 277     self.assertEqual(expected, actual)
 278
 279   def test_no_test_output(self):
 280     self.requests = [
 281       (
 282         '/get_result?r=key1',
 283         {'retry_404': False, 'retry_50x': False},
 284         generate_url_response(0, SWARM_OUTPUT_WITH_NO_TEST_OUTPUT, '0, 0'),
 285       ),
 286     ]
 287     expected = [gen_yielded_data(0, SWARM_OUTPUT_WITH_NO_TEST_OUTPUT, '0, 0')]
 288     actual = get_swarm_results(['key1'])
 289     self.assertEqual(expected, actual)
 290
 291   def test_no_keys(self):
 292     actual = get_swarm_results([])
 293     self.assertEqual([], actual)
 294
 295   def test_url_errors(self):
 296     self.mock(logging, 'error', lambda *_: None)
 297     # NOTE: get_swarm_results() hardcodes timeout=10. range(12) is because of an
 298     # additional time.time() call deep in net.url_open().
 299     now = {}
 300     lock = threading.Lock()
 301     def get_now():
 302       t = threading.current_thread()
 303       with lock:
 304         return now.setdefault(t, range(12)).pop(0)
 305     self.mock(swarming.net, 'sleep_before_retry', lambda _x, _y: None)
 306     self.mock(swarming, 'now', get_now)
 307     # The actual number of requests here depends on 'now' progressing to 10
 308     # seconds. It's called twice per loop.
 309     self.requests = [
 310       (
 311         '/get_result?r=key1',
 312         {'retry_404': False, 'retry_50x': False},
 313         None,
 314       ),
 315       (
 316         '/get_result?r=key1',
 317         {'retry_404': False, 'retry_50x': False},
 318         None,
 319       ),
 320       (
 321         '/get_result?r=key1',
 322         {'retry_404': False, 'retry_50x': False},
 323         None,
 324       ),
 325       (
 326         '/get_result?r=key1',
 327         {'retry_404': False, 'retry_50x': False},
 328         None,
 329       ),
 330       (
 331         '/get_result?r=key1',
 332         {'retry_404': False, 'retry_50x': False},
 333         None,
 334       ),
 335     ]
 336     actual = get_swarm_results(['key1'])
 337     self.assertEqual([], actual)
 338     self.assertTrue(all(not v for v in now.itervalues()), now)
 339
 340   def test_shard_repeated(self):
 341     self.requests = [
 342       (
 343         '/get_result?r=key1',
 344         {'retry_404': False, 'retry_50x': False},
 345         generate_url_response(0, SWARM_OUTPUT_SUCCESS, '0, 0'),
 346       ),
 347       (
 348         '/get_result?r=key1-repeat',
 349         {'retry_404': False, 'retry_50x': False},
 350         generate_url_response(0, SWARM_OUTPUT_SUCCESS, '0, 0'),
 351       ),
 352     ]
 353     expected = [gen_yielded_data(0, SWARM_OUTPUT_SUCCESS, '0, 0')]
 354     actual = get_swarm_results(['key1', 'key1-repeat'])
 355     self.assertEqual(expected, actual)
 356
 357   def test_one_shard_repeated(self):
 358     """Have shard 1 repeated twice, then shard 2 and 3."""
 359     self.requests = [
 360       (
 361         '/get_result?r=key1',
 362         {'retry_404': False, 'retry_50x': False},
 363         generate_url_response(0, TEST_SHARD_OUTPUT_1, '0, 0'),
 364       ),
 365       (
 366         '/get_result?r=key1-repeat',
 367         {'retry_404': False, 'retry_50x': False},
 368         generate_url_response(0, TEST_SHARD_OUTPUT_1, '0, 0'),
 369       ),
 370       (
 371         '/get_result?r=key2',
 372         {'retry_404': False, 'retry_50x': False},
 373         generate_url_response(1, TEST_SHARD_OUTPUT_2, '0, 0'),
 374       ),
 375       (
 376         '/get_result?r=key3',
 377         {'retry_404': False, 'retry_50x': False},
 378         generate_url_response(2, TEST_SHARD_OUTPUT_3, '0, 0'),
 379       ),
 380     ]
 381     expected = [
 382       gen_yielded_data(0, TEST_SHARD_OUTPUT_1, '0, 0'),
 383       gen_yielded_data(1, TEST_SHARD_OUTPUT_2, '0, 0'),
 384       gen_yielded_data(2, TEST_SHARD_OUTPUT_3, '0, 0'),
 385     ]
 386     actual = get_swarm_results(['key1', 'key1-repeat', 'key2', 'key3'])
 387     self.assertEqual(expected, sorted(actual))
 388
 389   def test_collect_nothing(self):
 390     self.mock(swarming, 'get_task_keys', lambda *_: [1, 2])
 391     self.mock(swarming, 'yield_results', lambda *_: [])
 392     self.assertEqual(
 393         1, swarming.collect('url', 'test_name', 'timeout', 'decorate'))
 394
 395   def test_collect_success(self):
 396     self.mock(swarming, 'get_task_keys', lambda *_: [1, 2])
 397     data = {
 398       'config_instance_index': 0,
 399       'exit_codes': '0',
 400       'machine_id': 0,
 401       'output': 'Foo',
 402     }
 403     self.mock(swarming, 'yield_results', lambda *_: [(0, data)])
 404     self.assertEqual(
 405         0, swarming.collect('url', 'test_name', 'timeout', 'decorate'))
 406     self._check_output(
 407         '\n================================================================\n'
 408         'Begin output from shard index 0 (machine tag: 0, id: unknown)\n'
 409         '================================================================\n\n'
 410         'Foo================================================================\n'
 411         'End output from shard index 0 (machine tag: 0, id: unknown). Return 0'
 412           '\n'
 413         '================================================================\n\n',
 414         '')
 415
 416   def test_collect_fail(self):
 417     self.mock(swarming, 'get_task_keys', lambda *_: [1, 2])
 418     data = {
 419       'config_instance_index': 0,
 420       'exit_codes': '0,8',
 421       'machine_id': 0,
 422       'output': 'Foo',
 423     }
 424     self.mock(swarming, 'yield_results', lambda *_: [(0, data)])
 425     self.assertEqual(
 426         8, swarming.collect('url', 'test_name', 'timeout', 'decorate'))
 427     self._check_output(
 428         '\n================================================================\n'
 429         'Begin output from shard index 0 (machine tag: 0, id: unknown)\n'
 430         '================================================================\n\n'
 431         'Foo================================================================\n'
 432         'End output from shard index 0 (machine tag: 0, id: unknown). Return 8'
 433           '\n'
 434         '================================================================\n\n',
 435         '')
 436
 437
 438 def chromium_tasks(retrieval_url, file_hash):
 439   return [
 440     {
 441       u'action': [
 442         u'python', u'run_isolated.zip',
 443         u'--hash', file_hash,
 444         u'--namespace', u'default-gzip',
 445         u'--isolate-server', retrieval_url,
 446       ],
 447       u'decorate_output': False,
 448       u'test_name': u'Run Test',
 449       u'time_out': 600,
 450     },
 451     {
 452       u'action' : [
 453           u'python', u'swarm_cleanup.py',
 454       ],
 455       u'decorate_output': False,
 456       u'test_name': u'Clean Up',
 457       u'time_out': 600,
 458     }
 459   ]
 460
 461
 462 def generate_expected_json(
 463     shards,
 464     dimensions,
 465     env,
 466     working_dir,
 467     isolate_server,
 468     profile,
 469     test_case_name=TEST_NAME,
 470     file_hash=FILE_HASH):
 471   expected = {
 472     u'cleanup': u'root',
 473     u'configurations': [
 474       {
 475         u'config_name': u'isolated',
 476         u'deadline_to_run': 60*60,
 477         u'dimensions': dimensions,
 478         u'min_instances': shards,
 479         u'priority': 101,
 480       },
 481     ],
 482     u'data': [],
 483     u'encoding': u'UTF-8',
 484     u'env_vars': env.copy(),
 485     u'restart_on_failure': True,
 486     u'test_case_name': test_case_name,
 487     u'tests': chromium_tasks(isolate_server, file_hash),
 488     u'working_dir': unicode(working_dir),
 489   }
 490   if shards > 1:
 491     expected[u'env_vars'][u'GTEST_SHARD_INDEX'] = u'%(instance_index)s'
 492     expected[u'env_vars'][u'GTEST_TOTAL_SHARDS'] = u'%(num_instances)s'
 493   if profile:
 494     expected[u'tests'][0][u'action'].append(u'--verbose')
 495   return expected
 496
 497
 498 class MockedStorage(object):
 499   def __init__(self, warm_cache):
 500     self._warm_cache = warm_cache
 501
 502   def __enter__(self):
 503     pass
 504
 505   def __exit__(self, *_args):
 506     pass
 507
 508   def upload_items(self, items):
 509     return [] if self._warm_cache else items
 510
 511   def get_fetch_url(self, _item):  # pylint: disable=R0201
 512     return 'http://localhost:8081/fetch_url'
 513
 514
 515 class ManifestTest(TestCase):
 516   def test_basic_manifest(self):
 517     env = {
 518       u'GTEST_SHARD_INDEX': u'%(instance_index)s',
 519       u'GTEST_TOTAL_SHARDS': u'%(num_instances)s',
 520     }
 521     dimensions = {'os': 'Windows'}
 522     manifest = swarming.Manifest(
 523         isolate_server='http://localhost:8081',
 524         namespace='default-gzip',
 525         isolated_hash=FILE_HASH,
 526         task_name=TEST_NAME,
 527         shards=2,
 528         env=env,
 529         dimensions=dimensions,
 530         working_dir='swarm_tests',
 531         deadline=60*60,
 532         verbose=False,
 533         profile=False,
 534         priority=101)
 535
 536     swarming.chromium_setup(manifest)
 537     manifest_json = json.loads(manifest.to_json())
 538
 539     expected = generate_expected_json(
 540         shards=2,
 541         dimensions={u'os': u'Windows'},
 542         env={},
 543         working_dir='swarm_tests',
 544         isolate_server=u'http://localhost:8081',
 545         profile=False)
 546     self.assertEqual(expected, manifest_json)
 547
 548   def test_basic_linux(self):
 549     """A basic linux manifest test to ensure that windows specific values
 550        aren't used.
 551     """
 552     dimensions = {'os': 'Linux'}
 553     manifest = swarming.Manifest(
 554         isolate_server='http://localhost:8081',
 555         namespace='default-gzip',
 556         isolated_hash=FILE_HASH,
 557         task_name=TEST_NAME,
 558         shards=1,
 559         env={},
 560         dimensions=dimensions,
 561         working_dir='swarm_tests',
 562         deadline=60*60,
 563         verbose=False,
 564         profile=False,
 565         priority=101)
 566
 567     swarming.chromium_setup(manifest)
 568     manifest_json = json.loads(manifest.to_json())
 569
 570     expected = generate_expected_json(
 571         shards=1,
 572         dimensions={u'os': u'Linux'},
 573         env={},
 574         working_dir='swarm_tests',
 575         isolate_server=u'http://localhost:8081',
 576         profile=False)
 577     self.assertEqual(expected, manifest_json)
 578
 579   def test_basic_linux_profile(self):
 580     dimensions = {'os': 'Linux'}
 581     manifest = swarming.Manifest(
 582         isolate_server='http://localhost:8081',
 583         namespace='default-gzip',
 584         isolated_hash=FILE_HASH,
 585         task_name=TEST_NAME,
 586         shards=1,
 587         env={},
 588         dimensions=dimensions,
 589         working_dir='swarm_tests',
 590         deadline=60*60,
 591         verbose=False,
 592         profile=True,
 593         priority=101)
 594
 595     swarming.chromium_setup(manifest)
 596     manifest_json = json.loads(manifest.to_json())
 597
 598     expected = generate_expected_json(
 599         shards=1,
 600         dimensions={u'os': u'Linux'},
 601         env={},
 602         working_dir='swarm_tests',
 603         isolate_server=u'http://localhost:8081',
 604         profile=True)
 605     self.assertEqual(expected, manifest_json)
 606
 607   def test_process_manifest_success(self):
 608     self.mock(swarming.net, 'url_read', lambda url, data=None: '{}')
 609     self.mock(swarming.isolateserver, 'get_storage',
 610         lambda *_: MockedStorage(warm_cache=False))
 611
 612     result = swarming.process_manifest(
 613         swarming='http://localhost:8082',
 614         isolate_server='http://localhost:8081',
 615         namespace='default',
 616         isolated_hash=FILE_HASH,
 617         task_name=TEST_NAME,
 618         shards=1,
 619         dimensions={},
 620         env={},
 621         working_dir='swarm_tests',
 622         deadline=60*60,
 623         verbose=False,
 624         profile=False,
 625         priority=101)
 626     self.assertEqual(0, result)
 627
 628   def test_process_manifest_success_zip_already_uploaded(self):
 629     self.mock(swarming.net, 'url_read', lambda url, data=None: '{}')
 630     self.mock(swarming.isolateserver, 'get_storage',
 631         lambda *_: MockedStorage(warm_cache=True))
 632
 633     dimensions = {'os': 'linux2'}
 634     result = swarming.process_manifest(
 635         swarming='http://localhost:8082',
 636         isolate_server='http://localhost:8081',
 637         namespace='default',
 638         isolated_hash=FILE_HASH,
 639         task_name=TEST_NAME,
 640         shards=1,
 641         dimensions=dimensions,
 642         env={},
 643         working_dir='swarm_tests',
 644         deadline=60*60,
 645         verbose=False,
 646         profile=False,
 647         priority=101)
 648     self.assertEqual(0, result)
 649
 650   def test_isolated_to_hash(self):
 651     calls = []
 652     self.mock(swarming.subprocess, 'call', lambda *c: calls.append(c))
 653     content = '{}'
 654     expected_hash = hashlib.sha1(content).hexdigest()
 655     handle, isolated = tempfile.mkstemp(
 656         prefix='swarming_test_', suffix='.isolated')
 657     os.close(handle)
 658     try:
 659       with open(isolated, 'w') as f:
 660         f.write(content)
 661       hash_value, is_file = swarming.isolated_to_hash(
 662           'http://localhost:1', 'default', isolated, hashlib.sha1, False)
 663     finally:
 664       os.remove(isolated)
 665     self.assertEqual(expected_hash, hash_value)
 666     self.assertEqual(True, is_file)
 667     expected_calls = [
 668         (
 669           [
 670             sys.executable,
 671             os.path.join(ROOT_DIR, 'isolate.py'),
 672             'archive',
 673             '--isolate-server', 'http://localhost:1',
 674             '--namespace', 'default',
 675             '--isolated',
 676             isolated,
 677           ],
 678           False,
 679         ),
 680     ]
 681     self.assertEqual(expected_calls, calls)
 682     self._check_output('Archiving: %s\n' % isolated, '')
 683
 684
 685 def mock_swarming_api_v1_bots():
 686   """Returns fake /swarming/api/v1/bots data."""
 687   now = datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')
 688   dead = '2006-01-02 03:04:05'
 689   return {
 690     'machine_death_timeout': 10,
 691     'machines': [
 692       {
 693         'tag': 'no-dimensions',
 694         'dimensions': {},
 695         'last_seen': now,
 696       },
 697       {
 698         'tag': 'amig1',
 699         'dimensions': {'os': 'amiga'},
 700         'last_seen': now,
 701       },
 702       {
 703         'tag': 'amig2',
 704         'dimensions': {'os': ['amiga', 'atari'], 'foo': 1},
 705         'last_seen': now,
 706       },
 707       {
 708         'tag': 'dead',
 709         'dimensions': {'os': 'amiga'},
 710         'last_seen': dead,
 711       },
 712     ],
 713   }
 714
 715
 716 class MainTest(TestCase):
 717   def setUp(self):
 718     super(MainTest, self).setUp()
 719     self._tmpdir = None
 720
 721   def tearDown(self):
 722     try:
 723       if self._tmpdir:
 724         shutil.rmtree(self._tmpdir)
 725     finally:
 726       super(MainTest, self).tearDown()
 727
 728   @property
 729   def tmpdir(self):
 730     if not self._tmpdir:
 731       self._tmpdir = tempfile.mkdtemp(prefix='swarming')
 732     return self._tmpdir
 733
 734   def test_run_hash(self):
 735     self.mock(swarming.isolateserver, 'get_storage',
 736         lambda *_: MockedStorage(warm_cache=False))
 737
 738     task_name = '%s/foo=bar_os=Mac/1111111111111111111111111111111111111111' % (
 739         getpass.getuser())
 740     j = generate_expected_json(
 741         shards=1,
 742         dimensions={'foo': 'bar', 'os': 'Mac'},
 743         env={},
 744         working_dir='swarm_tests',
 745         isolate_server='https://host2',
 746         profile=False,
 747         test_case_name=task_name)
 748     j['data'] = [['http://localhost:8081/fetch_url', 'swarm_data.zip']]
 749     data = {
 750       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
 751     }
 752     self.requests = [
 753       (
 754         '/test',
 755         {'data': data},
 756         # The actual output is ignored as long as it is valid json.
 757         StringIO.StringIO('{}'),
 758       ),
 759     ]
 760     ret = main([
 761         'trigger',
 762         '--swarming', 'https://host1',
 763         '--isolate-server', 'https://host2',
 764         '--shards', '1',
 765         '--priority', '101',
 766         '--dimension', 'foo', 'bar',
 767         '--dimension', 'os', 'Mac',
 768         '--deadline', '3600',
 769         FILE_HASH,
 770       ])
 771     actual = sys.stdout.getvalue()
 772     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
 773     self._check_output('Triggered task: %s\n' % task_name, '')
 774
 775   def test_run_isolated(self):
 776     self.mock(swarming.isolateserver, 'get_storage',
 777         lambda *_: MockedStorage(warm_cache=False))
 778     calls = []
 779     self.mock(swarming.subprocess, 'call', lambda *c: calls.append(c))
 780
 781     isolated = os.path.join(self.tmpdir, 'zaz.isolated')
 782     content = '{}'
 783     with open(isolated, 'wb') as f:
 784       f.write(content)
 785
 786     isolated_hash = ALGO(content).hexdigest()
 787     task_name = 'zaz/foo=bar_os=Mac/%s' % isolated_hash
 788     j = generate_expected_json(
 789         shards=1,
 790         dimensions={'foo': 'bar', 'os': 'Mac'},
 791         env={},
 792         working_dir='swarm_tests',
 793         isolate_server='https://host2',
 794         profile=False,
 795         test_case_name=task_name,
 796         file_hash=isolated_hash)
 797     j['data'] = [['http://localhost:8081/fetch_url', 'swarm_data.zip']]
 798     data = {
 799       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
 800     }
 801     self.requests = [
 802       (
 803         '/test',
 804         {'data': data},
 805         # The actual output is ignored as long as it is valid json.
 806         StringIO.StringIO('{}'),
 807       ),
 808     ]
 809     ret = main([
 810         'trigger',
 811         '--swarming', 'https://host1',
 812         '--isolate-server', 'https://host2',
 813         '--shards', '1',
 814         '--priority', '101',
 815         '--dimension', 'foo', 'bar',
 816         '--dimension', 'os', 'Mac',
 817         '--deadline', '3600',
 818         isolated,
 819       ])
 820     actual = sys.stdout.getvalue()
 821     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
 822     expected = [
 823       (
 824         [
 825           sys.executable,
 826           os.path.join(ROOT_DIR, 'isolate.py'), 'archive',
 827           '--isolate-server', 'https://host2',
 828           '--namespace' ,'default-gzip',
 829           '--isolated', isolated,
 830         ],
 831       0),
 832     ]
 833     self.assertEqual(expected, calls)
 834     expected = 'Archiving: %s\nTriggered task: %s\n' % (isolated, task_name)
 835     self._check_output(expected, '')
 836
 837   def test_query(self):
 838     self.requests = [
 839       (
 840         '/swarming/api/v1/bots',
 841         {
 842           "content_type": None,
 843           "data": None,
 844           "headers": None,
 845           "max_attempts": 30,
 846           "method": "GET",
 847           "read_timeout": 360.0,
 848           "retry_404": False,
 849           "retry_50x": True,
 850           "timeout": 360.0
 851         },
 852         StringIO.StringIO(json.dumps(mock_swarming_api_v1_bots())),
 853       ),
 854     ]
 855     main(['query', '--swarming', 'https://localhost:1'])
 856     expected = (
 857         "amig1\n  {u'os': u'amiga'}\n"
 858         "amig2\n  {u'foo': 1, u'os': [u'amiga', u'atari']}\n"
 859         "no-dimensions\n  {}\n")
 860     self._check_output(expected, '')
 861
 862   def test_query_bare(self):
 863     self.requests = [
 864       (
 865         '/swarming/api/v1/bots',
 866         {
 867           "content_type": None,
 868           "data": None,
 869           "headers": None,
 870           "max_attempts": 30,
 871           "method": "GET",
 872           "read_timeout": 360.0,
 873           "retry_404": False,
 874           "retry_50x": True,
 875           "timeout": 360.0
 876         },
 877         StringIO.StringIO(json.dumps(mock_swarming_api_v1_bots())),
 878       ),
 879     ]
 880     main(['query', '--swarming', 'https://localhost:1', '--bare'])
 881     self._check_output("amig1\namig2\nno-dimensions\n", '')
 882
 883   def test_query_filter(self):
 884     self.requests = [
 885       (
 886         '/swarming/api/v1/bots',
 887         {
 888           "content_type": None,
 889           "data": None,
 890           "headers": None,
 891           "max_attempts": 30,
 892           "method": "GET",
 893           "read_timeout": 360.0,
 894           "retry_404": False,
 895           "retry_50x": True,
 896           "timeout": 360.0
 897         },
 898         StringIO.StringIO(json.dumps(mock_swarming_api_v1_bots())),
 899       ),
 900     ]
 901     main(
 902         [
 903           'query', '--swarming', 'https://localhost:1',
 904           '--dimension', 'os', 'amiga',
 905         ])
 906     expected = (
 907         "amig1\n  {u'os': u'amiga'}\n"
 908         "amig2\n  {u'foo': 1, u'os': [u'amiga', u'atari']}\n")
 909     self._check_output(expected, '')
 910
 911   def test_query_filter_keep_dead(self):
 912     self.requests = [
 913       (
 914         '/swarming/api/v1/bots',
 915         {
 916           "content_type": None,
 917           "data": None,
 918           "headers": None,
 919           "max_attempts": 30,
 920           "method": "GET",
 921           "read_timeout": 360.0,
 922           "retry_404": False,
 923           "retry_50x": True,
 924           "timeout": 360.0
 925         },
 926         StringIO.StringIO(json.dumps(mock_swarming_api_v1_bots())),
 927       ),
 928     ]
 929     main(
 930         [
 931           'query', '--swarming', 'https://localhost:1',
 932           '--dimension', 'os', 'amiga', '--keep-dead',
 933         ])
 934     expected = (
 935         "amig1\n  {u'os': u'amiga'}\n"
 936         "amig2\n  {u'foo': 1, u'os': [u'amiga', u'atari']}\n"
 937         "dead\n  {u'os': u'amiga'}\n")
 938     self._check_output(expected, '')
 939
 940   def test_query_filter_dead_only(self):
 941     self.requests = [
 942       (
 943         '/swarming/api/v1/bots',
 944         {
 945           "content_type": None,
 946           "data": None,
 947           "headers": None,
 948           "max_attempts": 30,
 949           "method": "GET",
 950           "read_timeout": 360.0,
 951           "retry_404": False,
 952           "retry_50x": True,
 953           "timeout": 360.0
 954         },
 955         StringIO.StringIO(json.dumps(mock_swarming_api_v1_bots())),
 956       ),
 957     ]
 958     main(
 959         [
 960           'query', '--swarming', 'https://localhost:1',
 961           '--dimension', 'os', 'amiga', '--dead-only',
 962         ])
 963     expected = (
 964         "dead\n  {u'os': u'amiga'}\n")
 965     self._check_output(expected, '')
 966
 967   def test_trigger_no_request(self):
 968     with self.assertRaises(SystemExit):
 969       main([
 970             'trigger', '--swarming', 'https://host',
 971             '--isolate-server', 'https://host', '-T', 'foo',
 972           ])
 973     self._check_output(
 974         '',
 975         'Usage: swarming.py trigger [options] (hash|isolated)\n\n'
 976         'swarming.py: error: Must pass one .isolated file or its hash (sha1).'
 977         '\n')
 978
 979   def test_trigger_no_env_vars(self):
 980     with self.assertRaises(SystemExit):
 981       main(['trigger'])
 982     self._check_output(
 983         '',
 984         'Usage: swarming.py trigger [options] (hash|isolated)\n\n'
 985         'swarming.py: error: --swarming is required.'
 986         '\n')
 987
 988   def test_trigger_no_swarming_env_var(self):
 989     with self.assertRaises(SystemExit):
 990       with test_utils.EnvVars({'ISOLATE_SERVER': 'https://host'}):
 991         main(['trigger', '-T' 'foo', 'foo.isolated'])
 992     self._check_output(
 993         '',
 994         'Usage: swarming.py trigger [options] (hash|isolated)\n\n'
 995         'swarming.py: error: --swarming is required.'
 996         '\n')
 997
 998   def test_trigger_no_isolate_env_var(self):
 999     with self.assertRaises(SystemExit):
1000       with test_utils.EnvVars({'SWARMING_SERVER': 'https://host'}):
1001         main(['trigger', 'T', 'foo', 'foo.isolated'])
1002     self._check_output(
1003         '',
1004         'Usage: swarming.py trigger [options] (hash|isolated)\n\n'
1005         'swarming.py: error: Use one of --indir or --isolate-server.'
1006         '\n')
1007
1008   def test_trigger_env_var(self):
1009     with self.assertRaises(SystemExit):
1010       with test_utils.EnvVars({'ISOLATE_SERVER': 'https://host',
1011                                'SWARMING_SERVER': 'https://host'}):
1012         main(['trigger', '-T', 'foo'])
1013     self._check_output(
1014         '',
1015         'Usage: swarming.py trigger [options] (hash|isolated)\n\n'
1016         'swarming.py: error: Must pass one .isolated file or its hash (sha1).'
1017         '\n')
1018
1019   def test_trigger_no_task(self):
1020     with self.assertRaises(SystemExit):
1021       main([
1022             'trigger', '--swarming', 'https://host',
1023             '--isolate-server', 'https://host', 'foo.isolated',
1024           ])
1025     self._check_output(
1026         '',
1027         'Usage: swarming.py trigger [options] (hash|isolated)\n\n'
1028         'swarming.py: error: Please at least specify one --dimension\n')
1029
1030   def test_trigger_env(self):
1031     self.mock(swarming.isolateserver, 'get_storage',
1032         lambda *_: MockedStorage(warm_cache=False))
1033     j = generate_expected_json(
1034         shards=1,
1035         dimensions={'os': 'Mac'},
1036         env={'foo': 'bar'},
1037         working_dir='swarm_tests',
1038         isolate_server='https://host2',
1039         profile=False)
1040     j['data'] = [['http://localhost:8081/fetch_url', 'swarm_data.zip']]
1041     data = {
1042       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
1043     }
1044     self.requests = [
1045       (
1046         '/test',
1047         {'data': data},
1048         # The actual output is ignored as long as it is valid json.
1049         StringIO.StringIO('{}'),
1050       ),
1051     ]
1052     ret = main([
1053         'trigger',
1054         '--swarming', 'https://host1',
1055         '--isolate-server', 'https://host2',
1056         '--shards', '1',
1057         '--priority', '101',
1058         '--env', 'foo', 'bar',
1059         '--dimension', 'os', 'Mac',
1060         '--task-name', TEST_NAME,
1061         '--deadline', '3600',
1062         FILE_HASH,
1063       ])
1064     actual = sys.stdout.getvalue()
1065     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
1066
1067   def test_trigger_dimension_filter(self):
1068     self.mock(swarming.isolateserver, 'get_storage',
1069         lambda *_: MockedStorage(warm_cache=False))
1070     j = generate_expected_json(
1071         shards=1,
1072         dimensions={'foo': 'bar', 'os': 'Mac'},
1073         env={},
1074         working_dir='swarm_tests',
1075         isolate_server='https://host2',
1076         profile=False)
1077     j['data'] = [['http://localhost:8081/fetch_url', 'swarm_data.zip']]
1078     data = {
1079       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
1080     }
1081     self.requests = [
1082       (
1083         '/test',
1084         {'data': data},
1085         # The actual output is ignored as long as it is valid json.
1086         StringIO.StringIO('{}'),
1087       ),
1088     ]
1089     ret = main([
1090         'trigger',
1091         '--swarming', 'https://host1',
1092         '--isolate-server', 'https://host2',
1093         '--shards', '1',
1094         '--priority', '101',
1095         '--dimension', 'foo', 'bar',
1096         '--dimension', 'os', 'Mac',
1097         '--task-name', TEST_NAME,
1098         '--deadline', '3600',
1099         FILE_HASH,
1100       ])
1101     actual = sys.stdout.getvalue()
1102     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
1103
1104
1105 def clear_env_vars():
1106   for e in ('ISOLATE_SERVER', 'SWARMING_SERVER'):
1107     os.environ.pop(e, None)
1108
1109
1110 if __name__ == '__main__':
1111   logging.basicConfig(
1112       level=logging.DEBUG if '-v' in sys.argv else logging.ERROR)
1113   if '-v' in sys.argv:
1114     unittest.TestCase.maxDiff = None
1115   clear_env_vars()
1116   unittest.main()