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 inspect
  10 import json
  11 import logging
  12 import os
  13 import shutil
  14 import StringIO
  15 import sys
  16 import tempfile
  17 import threading
  18 import unittest
  19
  20 # net_utils adjusts sys.path.
  21 import net_utils
  22
  23 from depot_tools import auto_stub
  24
  25 import swarming
  26 import test_utils
  27
  28 from utils import net
  29 from utils import zip_package
  30
  31
  32 ALGO = hashlib.sha1
  33 FILE_NAME = u'test.isolated'
  34 FILE_HASH = u'1' * 40
  35 TEST_NAME = u'unit_tests'
  36
  37
  38 TEST_CASE_SUCCESS = (
  39   '[----------] 2 tests from StaticCookiePolicyTest\n'
  40   '[ RUN      ] StaticCookiePolicyTest.AllowAllCookiesTest\n'
  41   '[       OK ] StaticCookiePolicyTest.AllowAllCookiesTest (0 ms)\n'
  42   '[ RUN      ] StaticCookiePolicyTest.BlockAllCookiesTest\n'
  43   '[       OK ] StaticCookiePolicyTest.BlockAllCookiesTest (0 ms)\n'
  44   '[----------] 2 tests from StaticCookiePolicyTest (0 ms total)\n'
  45   '\n'
  46   '[----------] 1 test from TCPListenSocketTest\n'
  47   '[ RUN      ] TCPListenSocketTest.ServerSend\n'
  48   '[       OK ] TCPListenSocketTest.ServerSend (1 ms)\n'
  49   '[----------] 1 test from TCPListenSocketTest (1 ms total)\n')
  50
  51
  52 TEST_CASE_FAILURE = (
  53   '[----------] 2 tests from StaticCookiePolicyTest\n'
  54   '[ RUN      ] StaticCookiePolicyTest.AllowAllCookiesTest\n'
  55   '[       OK ] StaticCookiePolicyTest.AllowAllCookiesTest (0 ms)\n'
  56   '[ RUN      ] StaticCookiePolicyTest.BlockAllCookiesTest\n'
  57   'C:\\win\\build\\src\\chrome\\test.cc: error: Value of: result()\n'
  58   '  Actual: false\n'
  59   'Expected: true\n'
  60   '[  FAILED  ] StaticCookiePolicyTest.BlockAllCookiesTest (0 ms)\n'
  61   '[----------] 2 tests from StaticCookiePolicyTest (0 ms total)\n'
  62   '\n'
  63   '[----------] 1 test from TCPListenSocketTest\n'
  64   '[ RUN      ] TCPListenSocketTest.ServerSend\n'
  65   '[       OK ] TCPListenSocketTest.ServerSend (1 ms)\n'
  66   '[----------] 1 test from TCPListenSocketTest (1 ms total)\n')
  67
  68
  69 SWARM_OUTPUT_SUCCESS = (
  70   '[ RUN      ] unittests.Run Test\n' +
  71   TEST_CASE_SUCCESS +
  72   '[       OK ] unittests.Run Test (2549 ms)\n'
  73   '[ RUN      ] unittests.Clean Up\n'
  74   'No output!\n'
  75   '[       OK ] unittests.Clean Up (6 ms)\n'
  76   '\n'
  77   '[----------] unittests summary\n'
  78   '[==========] 2 tests ran. (2556 ms total)\n')
  79
  80
  81 SWARM_OUTPUT_FAILURE = (
  82   '[ RUN      ] unittests.Run Test\n' +
  83   TEST_CASE_FAILURE +
  84   '[       OK ] unittests.Run Test (2549 ms)\n'
  85   '[ RUN      ] unittests.Clean Up\n'
  86   'No output!\n'
  87   '[       OK ] unittests.Clean Up (6 ms)\n'
  88   '\n'
  89   '[----------] unittests summary\n'
  90   '[==========] 2 tests ran. (2556 ms total)\n')
  91
  92
  93 SWARM_OUTPUT_WITH_NO_TEST_OUTPUT = (
  94   '\n'
  95   'Unable to connection to swarm machine.\n')
  96
  97
  98 TEST_SHARD_1 = 'Note: This is test shard 1 of 3.'
  99 TEST_SHARD_2 = 'Note: This is test shard 2 of 3.'
 100 TEST_SHARD_3 = 'Note: This is test shard 3 of 3.'
 101
 102
 103 SWARM_SHARD_OUTPUT = (
 104   '[ RUN      ] unittests.Run Test\n'
 105   '%s\n'
 106   '[       OK ] unittests.Run Test (2549 ms)\n'
 107   '[ RUN      ] unittests.Clean Up\n'
 108   'No output!\n'
 109   '[       OK ] unittests.Clean Up (6 ms)\n'
 110   '\n'
 111   '[----------] unittests summary\n'
 112   '[==========] 2 tests ran. (2556 ms total)\n')
 113
 114
 115 TEST_SHARD_OUTPUT_1 = SWARM_SHARD_OUTPUT % TEST_SHARD_1
 116 TEST_SHARD_OUTPUT_2 = SWARM_SHARD_OUTPUT % TEST_SHARD_2
 117 TEST_SHARD_OUTPUT_3 = SWARM_SHARD_OUTPUT % TEST_SHARD_3
 118
 119 FAKE_BUNDLE_URL = 'http://localhost:8081/fetch_url'
 120
 121
 122 def gen_data(shard_output, exit_codes):
 123   return {
 124     u'config_instance_index': 0,
 125     u'exit_codes': unicode(exit_codes),
 126     u'machine_id': u'host',
 127     u'machine_tag': u'localhost',
 128     u'output': unicode(shard_output),
 129     u'isolated_out': swarming.extract_output_files_location(shard_output),
 130   }
 131
 132
 133 def gen_yielded_data(index, shard_output, exit_codes):
 134   """Returns an entry as it would be yielded by yield_results()."""
 135   return index, gen_data(shard_output, exit_codes)
 136
 137
 138 def generate_url_response(shard_output, exit_codes):
 139   return json.dumps(gen_data(shard_output, exit_codes))
 140
 141
 142 def get_swarm_results(keys, output_collector=None):
 143   """Simplifies the call to yield_results().
 144
 145   The timeout is hard-coded to 10 seconds.
 146   """
 147   return list(
 148       swarming.yield_results(
 149           'http://host:9001', keys, 10., None, True, output_collector))
 150
 151
 152 def collect(url, task_name, shards):
 153   """Simplifies the call to swarming.collect()."""
 154   return swarming.collect(
 155     url=url,
 156     task_name=task_name,
 157     shards=shards,
 158     timeout=10,
 159     decorate=True,
 160     print_status_updates=True,
 161     task_summary_json=None,
 162     task_output_dir=None)
 163
 164
 165 def gen_trigger_response(priority=101):
 166   # As seen in services/swarming/handlers_frontend.py.
 167   return {
 168     'priority': priority,
 169     'test_case_name': 'foo',
 170     'test_keys': [
 171       {
 172         'config_name': 'foo',
 173         'instance_index': 0,
 174         'num_instances': 1,
 175         'test_key': '123',
 176       }
 177     ],
 178   }
 179
 180
 181 def main(args):
 182   """Bypassies swarming.main()'s exception handling.
 183
 184   It gets in the way when debugging test failures.
 185   """
 186   dispatcher = swarming.subcommand.CommandDispatcher('swarming')
 187   return dispatcher.execute(swarming.OptionParserSwarming(), args)
 188
 189
 190 # Silence pylint 'Access to a protected member _Event of a client class'.
 191 class NonBlockingEvent(threading._Event):  # pylint: disable=W0212
 192   """Just like threading.Event, but a class and ignores timeout in 'wait'.
 193
 194   Intended to be used as a mock for threading.Event in tests.
 195   """
 196
 197   def wait(self, timeout=None):
 198     return super(NonBlockingEvent, self).wait(0)
 199
 200
 201 class TestCase(net_utils.TestCase):
 202   """Base class that defines the url_open mock."""
 203   def setUp(self):
 204     super(TestCase, self).setUp()
 205     self._lock = threading.Lock()
 206     self.mock(swarming.auth, 'ensure_logged_in', lambda _: None)
 207     self.mock(swarming.time, 'sleep', lambda _: None)
 208     self.mock(swarming.subprocess, 'call', lambda *_: self.fail())
 209     self.mock(swarming.threading, 'Event', NonBlockingEvent)
 210     self.mock(sys, 'stdout', StringIO.StringIO())
 211     self.mock(sys, 'stderr', StringIO.StringIO())
 212
 213   def tearDown(self):
 214     try:
 215       if not self.has_failed():
 216         self._check_output('', '')
 217     finally:
 218       super(TestCase, self).tearDown()
 219
 220   def _check_output(self, out, err):
 221     self.assertEqual(out, sys.stdout.getvalue())
 222     self.assertEqual(err, sys.stderr.getvalue())
 223
 224     # Flush their content by mocking them again.
 225     self.mock(sys, 'stdout', StringIO.StringIO())
 226     self.mock(sys, 'stderr', StringIO.StringIO())
 227
 228
 229 class TestGetTestKeys(TestCase):
 230   def test_no_keys(self):
 231     self.mock(swarming.time, 'sleep', lambda x: x)
 232     self.expected_requests(
 233         [
 234           (
 235             'http://host:9001/get_matching_test_cases?name=my_test',
 236             {'retry_404': True},
 237             'No matching Test Cases',
 238             None,
 239           ) for _ in range(net.URL_OPEN_MAX_ATTEMPTS)
 240         ])
 241     try:
 242       swarming.get_task_keys('http://host:9001', 'my_test')
 243       self.fail()
 244     except swarming.Failure as e:
 245       msg = (
 246           'Error: Unable to find any task with the name, my_test, on swarming '
 247           'server')
 248       self.assertEqual(msg, e.args[0])
 249
 250   def test_no_keys_on_first_attempt(self):
 251     self.mock(swarming.time, 'sleep', lambda x: x)
 252     keys = ['key_1', 'key_2']
 253     self.expected_requests(
 254         [
 255           (
 256             'http://host:9001/get_matching_test_cases?name=my_test',
 257             {'retry_404': True},
 258             'No matching Test Cases',
 259             None,
 260           ),
 261           (
 262             'http://host:9001/get_matching_test_cases?name=my_test',
 263             {'retry_404': True},
 264             json.dumps(keys),
 265             None,
 266           ),
 267         ])
 268     actual = swarming.get_task_keys('http://host:9001', 'my_test')
 269     self.assertEqual(keys, actual)
 270
 271   def test_find_keys(self):
 272     keys = ['key_1', 'key_2']
 273     self.expected_requests(
 274         [
 275           (
 276             'http://host:9001/get_matching_test_cases?name=my_test',
 277             {'retry_404': True},
 278             json.dumps(keys),
 279             None,
 280           ),
 281         ])
 282     actual = swarming.get_task_keys('http://host:9001', 'my_test')
 283     self.assertEqual(keys, actual)
 284
 285
 286 class TestGetSwarmResults(TestCase):
 287   def test_success(self):
 288     self.expected_requests(
 289         [
 290           (
 291             'http://host:9001/get_result?r=key1',
 292             {'retry_404': False, 'retry_50x': False},
 293             generate_url_response(SWARM_OUTPUT_SUCCESS, '0, 0'),
 294             None,
 295           ),
 296         ])
 297     expected = [gen_yielded_data(0, SWARM_OUTPUT_SUCCESS, '0, 0')]
 298     actual = get_swarm_results(['key1'])
 299     self.assertEqual(expected, actual)
 300
 301   def test_failure(self):
 302     self.expected_requests(
 303         [
 304           (
 305             'http://host:9001/get_result?r=key1',
 306             {'retry_404': False, 'retry_50x': False},
 307             generate_url_response(SWARM_OUTPUT_FAILURE, '0, 1'),
 308             None,
 309           ),
 310         ])
 311     expected = [gen_yielded_data(0, SWARM_OUTPUT_FAILURE, '0, 1')]
 312     actual = get_swarm_results(['key1'])
 313     self.assertEqual(expected, actual)
 314
 315   def test_no_test_output(self):
 316     self.expected_requests(
 317         [
 318           (
 319             'http://host:9001/get_result?r=key1',
 320             {'retry_404': False, 'retry_50x': False},
 321             generate_url_response(SWARM_OUTPUT_WITH_NO_TEST_OUTPUT, '0, 0'),
 322             None,
 323           ),
 324         ])
 325     expected = [gen_yielded_data(0, SWARM_OUTPUT_WITH_NO_TEST_OUTPUT, '0, 0')]
 326     actual = get_swarm_results(['key1'])
 327     self.assertEqual(expected, actual)
 328
 329   def test_no_keys(self):
 330     actual = get_swarm_results([])
 331     self.assertEqual([], actual)
 332
 333   def test_url_errors(self):
 334     self.mock(logging, 'error', lambda *_, **__: None)
 335     # NOTE: get_swarm_results() hardcodes timeout=10.
 336     now = {}
 337     lock = threading.Lock()
 338     def get_now():
 339       t = threading.current_thread()
 340       with lock:
 341         return now.setdefault(t, range(10)).pop(0)
 342     self.mock(swarming.net, 'sleep_before_retry', lambda _x, _y: None)
 343     self.mock(swarming, 'now', get_now)
 344     # The actual number of requests here depends on 'now' progressing to 10
 345     # seconds. It's called once per loop. Loop makes 9 iterations.
 346     self.expected_requests(
 347         9 * [
 348           (
 349             'http://host:9001/get_result?r=key1',
 350             {'retry_404': False, 'retry_50x': False},
 351             None,
 352             None,
 353           )
 354         ])
 355     actual = get_swarm_results(['key1'])
 356     self.assertEqual([], actual)
 357     self.assertTrue(all(not v for v in now.itervalues()), now)
 358
 359   def test_many_shards(self):
 360     self.expected_requests(
 361         [
 362           (
 363             'http://host:9001/get_result?r=key1',
 364             {'retry_404': False, 'retry_50x': False},
 365             generate_url_response(TEST_SHARD_OUTPUT_1, '0, 0'),
 366             None,
 367           ),
 368           (
 369             'http://host:9001/get_result?r=key2',
 370             {'retry_404': False, 'retry_50x': False},
 371             generate_url_response(TEST_SHARD_OUTPUT_2, '0, 0'),
 372             None,
 373           ),
 374           (
 375             'http://host:9001/get_result?r=key3',
 376             {'retry_404': False, 'retry_50x': False},
 377             generate_url_response(TEST_SHARD_OUTPUT_3, '0, 0'),
 378             None,
 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', 'key2', 'key3'])
 387     self.assertEqual(expected, sorted(actual))
 388
 389   def test_output_collector_called(self):
 390     # Three shards, one failed. All results are passed to output collector.
 391     self.expected_requests(
 392         [
 393           (
 394             'http://host:9001/get_result?r=key1',
 395             {'retry_404': False, 'retry_50x': False},
 396             generate_url_response(TEST_SHARD_OUTPUT_1, '0, 0'),
 397             None,
 398           ),
 399           (
 400             'http://host:9001/get_result?r=key2',
 401             {'retry_404': False, 'retry_50x': False},
 402             generate_url_response(TEST_SHARD_OUTPUT_2, '0, 0'),
 403             None,
 404           ),
 405           (
 406             'http://host:9001/get_result?r=key3',
 407             {'retry_404': False, 'retry_50x': False},
 408             generate_url_response(SWARM_OUTPUT_FAILURE, '0, 1'),
 409             None,
 410           ),
 411         ])
 412
 413     class FakeOutputCollector(object):
 414       def __init__(self):
 415         self.results = []
 416         self._lock = threading.Lock()
 417
 418       def process_shard_result(self, index, result):
 419         with self._lock:
 420           self.results.append((index, result))
 421
 422     output_collector = FakeOutputCollector()
 423     get_swarm_results(['key1', 'key2', 'key3'], output_collector)
 424
 425     expected = [
 426       (0, gen_data(TEST_SHARD_OUTPUT_1, '0, 0')),
 427       (1, gen_data(TEST_SHARD_OUTPUT_2, '0, 0')),
 428       (2, gen_data(SWARM_OUTPUT_FAILURE, '0, 1')),
 429     ]
 430     self.assertEqual(sorted(expected), sorted(output_collector.results))
 431
 432   def test_collect_nothing(self):
 433     self.mock(swarming, 'get_task_keys', lambda *_: ['task_key'])
 434     self.mock(swarming, 'yield_results', lambda *_: [])
 435     self.assertEqual(1, collect('url', 'name', 2))
 436     self._check_output('', 'Results from some shards are missing: 0, 1\n')
 437
 438   def test_collect_success(self):
 439     self.mock(swarming, 'get_task_keys', lambda *_: ['task_key'])
 440     data = {
 441       'config_instance_index': 0,
 442       'exit_codes': '0',
 443       'machine_id': 0,
 444       'output': 'Foo\n',
 445     }
 446     self.mock(swarming, 'yield_results', lambda *_: [(0, data)])
 447     self.assertEqual(0, collect('url', 'name', 1))
 448     self._check_output(
 449         '\n================================================================\n'
 450         'Begin output from shard index 0 (machine tag: 0, id: unknown)\n'
 451         '================================================================\n\n'
 452         'Foo\n'
 453         '================================================================\n'
 454         'End output from shard index 0 (machine tag: 0, id: unknown).\n'
 455         'Exit code 0 (0x0).\n'
 456         '================================================================\n\n',
 457         '')
 458
 459   def test_collect_fail(self):
 460     self.mock(swarming, 'get_task_keys', lambda *_: ['task_key'])
 461     data = {
 462       'config_instance_index': 0,
 463       'exit_codes': '0,8',
 464       'machine_id': 0,
 465       'output': 'Foo\n',
 466     }
 467     self.mock(swarming, 'yield_results', lambda *_: [(0, data)])
 468     self.assertEqual(1, collect('url', 'name', 1))
 469     self._check_output(
 470         '\n================================================================\n'
 471         'Begin output from shard index 0 (machine tag: 0, id: unknown)\n'
 472         '================================================================\n\n'
 473         'Foo\n'
 474         '================================================================\n'
 475         'End output from shard index 0 (machine tag: 0, id: unknown).\n'
 476         'Exit code 8 (0x8).\n'
 477         '================================================================\n\n',
 478         '')
 479
 480   def test_collect_negative_exit_code(self):
 481     self.mock(swarming, 'get_task_keys', lambda *_: ['task_key'])
 482     data = {
 483       'config_instance_index': 0,
 484       'exit_codes': '-1073741515,0',
 485       'machine_id': 0,
 486       'output': 'Foo\n',
 487     }
 488     self.mock(swarming, 'yield_results', lambda *_: [(0, data)])
 489     self.assertEqual(1, collect('url', 'name', 1))
 490     self._check_output(
 491         '\n================================================================\n'
 492         'Begin output from shard index 0 (machine tag: 0, id: unknown)\n'
 493         '================================================================\n\n'
 494         'Foo\n'
 495         '================================================================\n'
 496         'End output from shard index 0 (machine tag: 0, id: unknown).\n'
 497         'Exit code -1073741515 (0xc0000135).\n'
 498         '================================================================\n\n',
 499         '')
 500
 501   def test_collect_one_missing(self):
 502     self.mock(swarming, 'get_task_keys', lambda *_: ['task_key'])
 503     data = {
 504       'config_instance_index': 0,
 505       'exit_codes': '0',
 506       'machine_id': 0,
 507       'output': 'Foo\n',
 508     }
 509     self.mock(swarming, 'yield_results', lambda *_: [(0, data)])
 510     self.assertEqual(1, collect('url', 'name', 2))
 511     self._check_output(
 512         '\n================================================================\n'
 513         'Begin output from shard index 0 (machine tag: 0, id: unknown)\n'
 514         '================================================================\n\n'
 515         'Foo\n'
 516         '================================================================\n'
 517         'End output from shard index 0 (machine tag: 0, id: unknown).\n'
 518         'Exit code 0 (0x0).\n'
 519         '================================================================\n\n',
 520         'Results from some shards are missing: 1\n')
 521
 522
 523 def chromium_tasks(retrieval_url, file_hash, extra_args):
 524   return [
 525     {
 526       u'action': [
 527         u'python', u'run_isolated.zip',
 528         u'--hash', file_hash,
 529         u'--namespace', u'default-gzip',
 530         u'--isolate-server', retrieval_url,
 531       ] + (['--'] + list(extra_args) if extra_args else []),
 532       u'decorate_output': False,
 533       u'test_name': u'Run Test',
 534       u'hard_time_out': 2*60*60,
 535     },
 536     {
 537       u'action' : [
 538           u'python', u'swarm_cleanup.py',
 539       ],
 540       u'decorate_output': False,
 541       u'test_name': u'Clean Up',
 542       u'hard_time_out': 2*60*60,
 543     }
 544   ]
 545
 546
 547 def generate_expected_json(
 548     shards,
 549     shard_index,
 550     dimensions,
 551     env,
 552     isolate_server,
 553     profile,
 554     test_case_name=TEST_NAME,
 555     file_hash=FILE_HASH,
 556     extra_args=None):
 557   expected = {
 558     u'cleanup': u'root',
 559     u'configurations': [
 560       {
 561         u'config_name': u'isolated',
 562         u'deadline_to_run': 60*60,
 563         u'dimensions': dimensions,
 564         u'priority': 101,
 565       },
 566     ],
 567     u'data': [],
 568     u'env_vars': env.copy(),
 569     u'test_case_name': test_case_name,
 570     u'tests': chromium_tasks(isolate_server, file_hash, extra_args),
 571   }
 572   if shards > 1:
 573     expected[u'env_vars'][u'GTEST_SHARD_INDEX'] = u'%d' % shard_index
 574     expected[u'env_vars'][u'GTEST_TOTAL_SHARDS'] = u'%d' % shards
 575   if profile:
 576     expected[u'tests'][0][u'action'].append(u'--verbose')
 577   return expected
 578
 579
 580 class MockedStorage(object):
 581   def __init__(self, warm_cache):
 582     self._warm_cache = warm_cache
 583
 584   def __enter__(self):
 585     return self
 586
 587   def __exit__(self, *_args):
 588     pass
 589
 590   def upload_items(self, items):
 591     return [] if self._warm_cache else items
 592
 593   def get_fetch_url(self, _item):  # pylint: disable=R0201
 594     return FAKE_BUNDLE_URL
 595
 596
 597 class TriggerTaskShardsTest(TestCase):
 598   def test_zip_bundle_files(self):
 599     manifest = swarming.Manifest(
 600         isolate_server='http://localhost:8081',
 601         namespace='default-gzip',
 602         isolated_hash=FILE_HASH,
 603         task_name=TEST_NAME,
 604         extra_args=None,
 605         env={},
 606         dimensions={'os': 'Linux'},
 607         deadline=60*60,
 608         verbose=False,
 609         profile=False,
 610         priority=101)
 611
 612     bundle = zip_package.ZipPackage(swarming.ROOT_DIR)
 613     swarming.setup_run_isolated(manifest, bundle)
 614
 615     self.assertEqual(
 616         set(['run_isolated.zip', 'swarm_cleanup.py']), set(bundle.files))
 617
 618   def test_basic(self):
 619     manifest = swarming.Manifest(
 620         isolate_server='http://localhost:8081',
 621         namespace='default-gzip',
 622         isolated_hash=FILE_HASH,
 623         task_name=TEST_NAME,
 624         extra_args=None,
 625         env={},
 626         dimensions={'os': 'Linux'},
 627         deadline=60*60,
 628         verbose=False,
 629         profile=False,
 630         priority=101)
 631
 632     swarming.setup_run_isolated(manifest, None)
 633     manifest_json = json.loads(manifest.to_json())
 634
 635     expected = generate_expected_json(
 636         shards=1,
 637         shard_index=0,
 638         dimensions={u'os': u'Linux'},
 639         env={},
 640         isolate_server=u'http://localhost:8081',
 641         profile=False)
 642     self.assertEqual(expected, manifest_json)
 643
 644   def test_basic_profile(self):
 645     manifest = swarming.Manifest(
 646         isolate_server='http://localhost:8081',
 647         namespace='default-gzip',
 648         isolated_hash=FILE_HASH,
 649         task_name=TEST_NAME,
 650         extra_args=None,
 651         env={},
 652         dimensions={'os': 'Linux'},
 653         deadline=60*60,
 654         verbose=False,
 655         profile=True,
 656         priority=101)
 657
 658     swarming.setup_run_isolated(manifest, None)
 659     manifest_json = json.loads(manifest.to_json())
 660
 661     expected = generate_expected_json(
 662         shards=1,
 663         shard_index=0,
 664         dimensions={u'os': u'Linux'},
 665         env={},
 666         isolate_server=u'http://localhost:8081',
 667         profile=True)
 668     self.assertEqual(expected, manifest_json)
 669
 670   def test_manifest_with_extra_args(self):
 671     manifest = swarming.Manifest(
 672         isolate_server='http://localhost:8081',
 673         namespace='default-gzip',
 674         isolated_hash=FILE_HASH,
 675         task_name=TEST_NAME,
 676         extra_args=['--extra-cmd-arg=1234', 'some more'],
 677         env={},
 678         dimensions={'os': 'Windows'},
 679         deadline=60*60,
 680         verbose=False,
 681         profile=False,
 682         priority=101)
 683
 684     swarming.setup_run_isolated(manifest, None)
 685     manifest_json = json.loads(manifest.to_json())
 686
 687     expected = generate_expected_json(
 688         shards=1,
 689         shard_index=0,
 690         dimensions={u'os': u'Windows'},
 691         env={},
 692         isolate_server=u'http://localhost:8081',
 693         profile=False,
 694         extra_args=['--extra-cmd-arg=1234', 'some more'])
 695     self.assertEqual(expected, manifest_json)
 696
 697   def test_manifest_for_shard(self):
 698     manifest = swarming.Manifest(
 699         isolate_server='http://localhost:8081',
 700         namespace='default-gzip',
 701         isolated_hash=FILE_HASH,
 702         task_name=TEST_NAME,
 703         extra_args=None,
 704         env=swarming.setup_googletest({}, 5, 3),
 705         dimensions={'os': 'Linux'},
 706         deadline=60*60,
 707         verbose=False,
 708         profile=False,
 709         priority=101)
 710
 711     swarming.setup_run_isolated(manifest, None)
 712     manifest_json = json.loads(manifest.to_json())
 713
 714     expected = generate_expected_json(
 715         shards=5,
 716         shard_index=3,
 717         dimensions={u'os': u'Linux'},
 718         env={},
 719         isolate_server=u'http://localhost:8081',
 720         profile=False)
 721     self.assertEqual(expected, manifest_json)
 722
 723   def test_trigger_task_shards_success(self):
 724     self.mock(
 725         swarming.net, 'url_read',
 726         lambda url, data=None: json.dumps(gen_trigger_response()))
 727     self.mock(swarming.isolateserver, 'get_storage',
 728         lambda *_: MockedStorage(warm_cache=False))
 729
 730     tasks = swarming.trigger_task_shards(
 731         swarming='http://localhost:8082',
 732         isolate_server='http://localhost:8081',
 733         namespace='default',
 734         isolated_hash=FILE_HASH,
 735         task_name=TEST_NAME,
 736         extra_args=['--some-arg', '123'],
 737         shards=1,
 738         dimensions={},
 739         env={},
 740         deadline=60*60,
 741         verbose=False,
 742         profile=False,
 743         priority=101)
 744     expected = {
 745       'unit_tests': {
 746         'shard_index': 0,
 747         'task_id': '123',
 748         'view_url': 'http://localhost:8082/user/task/123',
 749       }
 750     }
 751     self.assertEqual(expected, tasks)
 752
 753   def test_trigger_task_shards_priority_override(self):
 754     self.mock(
 755         swarming.net, 'url_read',
 756         lambda url, data=None: json.dumps(gen_trigger_response(priority=200)))
 757     self.mock(swarming.isolateserver, 'get_storage',
 758         lambda *_: MockedStorage(warm_cache=False))
 759
 760     tasks = swarming.trigger_task_shards(
 761         swarming='http://localhost:8082',
 762         isolate_server='http://localhost:8081',
 763         namespace='default',
 764         isolated_hash=FILE_HASH,
 765         task_name=TEST_NAME,
 766         extra_args=['--some-arg', '123'],
 767         shards=2,
 768         dimensions={},
 769         env={},
 770         deadline=60*60,
 771         verbose=False,
 772         profile=False,
 773         priority=101)
 774     expected = {
 775       u'unit_tests:2:0': {
 776         u'shard_index': 0,
 777         u'task_id': u'123',
 778         u'view_url': u'http://localhost:8082/user/task/123',
 779       },
 780       u'unit_tests:2:1': {
 781         u'shard_index': 1,
 782         u'task_id': u'123',
 783         u'view_url': u'http://localhost:8082/user/task/123',
 784       }
 785     }
 786     self.assertEqual(expected, tasks)
 787     self._check_output('', 'Priority was reset to 200\n')
 788
 789   def test_trigger_task_shards_success_zip_already_uploaded(self):
 790     self.mock(
 791         swarming.net, 'url_read',
 792         lambda url, data=None: json.dumps(gen_trigger_response()))
 793     self.mock(swarming.isolateserver, 'get_storage',
 794         lambda *_: MockedStorage(warm_cache=True))
 795
 796     dimensions = {'os': 'linux2'}
 797     tasks = swarming.trigger_task_shards(
 798         swarming='http://localhost:8082',
 799         isolate_server='http://localhost:8081',
 800         namespace='default',
 801         isolated_hash=FILE_HASH,
 802         task_name=TEST_NAME,
 803         extra_args=['--some-arg', '123'],
 804         shards=1,
 805         dimensions=dimensions,
 806         env={},
 807         deadline=60*60,
 808         verbose=False,
 809         profile=False,
 810         priority=101)
 811
 812     expected = {
 813       'unit_tests': {
 814         'shard_index': 0,
 815         'task_id': '123',
 816         'view_url': 'http://localhost:8082/user/task/123',
 817       }
 818     }
 819     self.assertEqual(expected, tasks)
 820
 821   def test_isolated_to_hash(self):
 822     calls = []
 823     self.mock(swarming.subprocess, 'call', lambda *c: calls.append(c))
 824     content = '{}'
 825     expected_hash = hashlib.sha1(content).hexdigest()
 826     handle, isolated = tempfile.mkstemp(
 827         prefix='swarming_test_', suffix='.isolated')
 828     os.close(handle)
 829     try:
 830       with open(isolated, 'w') as f:
 831         f.write(content)
 832       hash_value, is_file = swarming.isolated_to_hash(
 833           'http://localhost:1', 'default', isolated, hashlib.sha1, False)
 834     finally:
 835       os.remove(isolated)
 836     self.assertEqual(expected_hash, hash_value)
 837     self.assertEqual(True, is_file)
 838     expected_calls = [
 839         (
 840           [
 841             sys.executable,
 842             os.path.join(swarming.ROOT_DIR, 'isolate.py'),
 843             'archive',
 844             '--isolate-server', 'http://localhost:1',
 845             '--namespace', 'default',
 846             '--isolated',
 847             isolated,
 848           ],
 849           False,
 850         ),
 851     ]
 852     self.assertEqual(expected_calls, calls)
 853     self._check_output('Archiving: %s\n' % isolated, '')
 854
 855
 856 class MainTest(TestCase):
 857   def setUp(self):
 858     super(MainTest, self).setUp()
 859     self._tmpdir = None
 860
 861   def tearDown(self):
 862     try:
 863       if self._tmpdir:
 864         shutil.rmtree(self._tmpdir)
 865     finally:
 866       super(MainTest, self).tearDown()
 867
 868   @property
 869   def tmpdir(self):
 870     if not self._tmpdir:
 871       self._tmpdir = tempfile.mkdtemp(prefix='swarming')
 872     return self._tmpdir
 873
 874   def test_run_hash(self):
 875     self.mock(swarming.isolateserver, 'get_storage',
 876         lambda *_: MockedStorage(warm_cache=False))
 877     self.mock(swarming, 'now', lambda: 123456)
 878
 879     task_name = (
 880         '%s/foo=bar_os=Mac/1111111111111111111111111111111111111111/123456000' %
 881         getpass.getuser())
 882     j = generate_expected_json(
 883         shards=1,
 884         shard_index=0,
 885         dimensions={'foo': 'bar', 'os': 'Mac'},
 886         env={},
 887         isolate_server='https://host2',
 888         profile=False,
 889         test_case_name=task_name)
 890     j['data'] = [[FAKE_BUNDLE_URL, 'swarm_data.zip']]
 891     data = {
 892       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
 893     }
 894     self.expected_requests(
 895         [
 896           (
 897             'https://host1/test',
 898             {'data': data},
 899             json.dumps(gen_trigger_response()),
 900             None,
 901           ),
 902         ])
 903     ret = main([
 904         'trigger',
 905         '--swarming', 'https://host1',
 906         '--isolate-server', 'https://host2',
 907         '--shards', '1',
 908         '--priority', '101',
 909         '--dimension', 'foo', 'bar',
 910         '--dimension', 'os', 'Mac',
 911         '--deadline', '3600',
 912         FILE_HASH,
 913       ])
 914     actual = sys.stdout.getvalue()
 915     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
 916     self._check_output('Triggered task: %s\n' % task_name, '')
 917
 918   def test_run_isolated(self):
 919     self.mock(swarming.isolateserver, 'get_storage',
 920         lambda *_: MockedStorage(warm_cache=False))
 921     calls = []
 922     self.mock(swarming.subprocess, 'call', lambda *c: calls.append(c))
 923     self.mock(swarming, 'now', lambda: 123456)
 924
 925     isolated = os.path.join(self.tmpdir, 'zaz.isolated')
 926     content = '{}'
 927     with open(isolated, 'wb') as f:
 928       f.write(content)
 929
 930     isolated_hash = ALGO(content).hexdigest()
 931     task_name = 'zaz/foo=bar_os=Mac/%s/123456000' % isolated_hash
 932     j = generate_expected_json(
 933         shards=1,
 934         shard_index=0,
 935         dimensions={'foo': 'bar', 'os': 'Mac'},
 936         env={},
 937         isolate_server='https://host2',
 938         profile=False,
 939         test_case_name=task_name,
 940         file_hash=isolated_hash)
 941     j['data'] = [[FAKE_BUNDLE_URL, 'swarm_data.zip']]
 942     data = {
 943       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
 944     }
 945     self.expected_requests(
 946         [
 947           (
 948             'https://host1/test',
 949             {'data': data},
 950             json.dumps(gen_trigger_response()),
 951             None,
 952           ),
 953         ])
 954     ret = main([
 955         'trigger',
 956         '--swarming', 'https://host1',
 957         '--isolate-server', 'https://host2',
 958         '--shards', '1',
 959         '--priority', '101',
 960         '--dimension', 'foo', 'bar',
 961         '--dimension', 'os', 'Mac',
 962         '--deadline', '3600',
 963         isolated,
 964       ])
 965     actual = sys.stdout.getvalue()
 966     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
 967     expected = [
 968       (
 969         [
 970           sys.executable,
 971           os.path.join(swarming.ROOT_DIR, 'isolate.py'), 'archive',
 972           '--isolate-server', 'https://host2',
 973           '--namespace' ,'default-gzip',
 974           '--isolated', isolated,
 975         ],
 976       0),
 977     ]
 978     self.assertEqual(expected, calls)
 979     self._check_output(
 980         'Archiving: %s\nTriggered task: %s\n' % (isolated, task_name), '')
 981
 982   def test_trigger_no_request(self):
 983     with self.assertRaises(SystemExit):
 984       main([
 985             'trigger', '--swarming', 'https://host',
 986             '--isolate-server', 'https://host', '-T', 'foo',
 987           ])
 988     self._check_output(
 989         '',
 990         'Usage: swarming.py trigger [options] (hash|isolated) [-- extra_args]'
 991         '\n\n'
 992         'swarming.py: error: Must pass one .isolated file or its hash (sha1).'
 993         '\n')
 994
 995   def test_trigger_no_env_vars(self):
 996     with self.assertRaises(SystemExit):
 997       main(['trigger'])
 998     self._check_output(
 999         '',
1000         'Usage: swarming.py trigger [options] (hash|isolated) [-- extra_args]'
1001         '\n\n'
1002         'swarming.py: error: --swarming is required.'
1003         '\n')
1004
1005   def test_trigger_no_swarming_env_var(self):
1006     with self.assertRaises(SystemExit):
1007       with test_utils.EnvVars({'ISOLATE_SERVER': 'https://host'}):
1008         main(['trigger', '-T' 'foo', 'foo.isolated'])
1009     self._check_output(
1010         '',
1011         'Usage: swarming.py trigger [options] (hash|isolated) [-- extra_args]'
1012         '\n\n'
1013         'swarming.py: error: --swarming is required.'
1014         '\n')
1015
1016   def test_trigger_no_isolate_env_var(self):
1017     with self.assertRaises(SystemExit):
1018       with test_utils.EnvVars({'SWARMING_SERVER': 'https://host'}):
1019         main(['trigger', 'T', 'foo', 'foo.isolated'])
1020     self._check_output(
1021         '',
1022         'Usage: swarming.py trigger [options] (hash|isolated) [-- extra_args]'
1023         '\n\n'
1024         'swarming.py: error: Use one of --indir or --isolate-server.'
1025         '\n')
1026
1027   def test_trigger_env_var(self):
1028     with self.assertRaises(SystemExit):
1029       with test_utils.EnvVars({'ISOLATE_SERVER': 'https://host',
1030                                'SWARMING_SERVER': 'https://host'}):
1031         main(['trigger', '-T', 'foo'])
1032     self._check_output(
1033         '',
1034         'Usage: swarming.py trigger [options] (hash|isolated) [-- extra_args]'
1035         '\n\n'
1036         'swarming.py: error: Must pass one .isolated file or its hash (sha1).'
1037         '\n')
1038
1039   def test_trigger_no_task(self):
1040     with self.assertRaises(SystemExit):
1041       main([
1042             'trigger', '--swarming', 'https://host',
1043             '--isolate-server', 'https://host', 'foo.isolated',
1044           ])
1045     self._check_output(
1046         '',
1047         'Usage: swarming.py trigger [options] (hash|isolated) [-- extra_args]'
1048         '\n\n'
1049         'swarming.py: error: Please at least specify one --dimension\n')
1050
1051   def test_trigger_env(self):
1052     self.mock(swarming.isolateserver, 'get_storage',
1053         lambda *_: MockedStorage(warm_cache=False))
1054     j = generate_expected_json(
1055         shards=1,
1056         shard_index=0,
1057         dimensions={'os': 'Mac'},
1058         env={'foo': 'bar'},
1059         isolate_server='https://host2',
1060         profile=False)
1061     j['data'] = [[FAKE_BUNDLE_URL, 'swarm_data.zip']]
1062     data = {
1063       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
1064     }
1065     self.expected_requests(
1066         [
1067           (
1068             'https://host1/test',
1069             {'data': data},
1070             json.dumps(gen_trigger_response()),
1071             None,
1072           ),
1073         ])
1074     ret = main([
1075         'trigger',
1076         '--swarming', 'https://host1',
1077         '--isolate-server', 'https://host2',
1078         '--shards', '1',
1079         '--priority', '101',
1080         '--env', 'foo', 'bar',
1081         '--dimension', 'os', 'Mac',
1082         '--task-name', TEST_NAME,
1083         '--deadline', '3600',
1084         FILE_HASH,
1085       ])
1086     actual = sys.stdout.getvalue()
1087     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
1088
1089   def test_trigger_dimension_filter(self):
1090     self.mock(swarming.isolateserver, 'get_storage',
1091         lambda *_: MockedStorage(warm_cache=False))
1092     j = generate_expected_json(
1093         shards=1,
1094         shard_index=0,
1095         dimensions={'foo': 'bar', 'os': 'Mac'},
1096         env={},
1097         isolate_server='https://host2',
1098         profile=False)
1099     j['data'] = [[FAKE_BUNDLE_URL, 'swarm_data.zip']]
1100     data = {
1101       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
1102     }
1103     self.expected_requests(
1104         [
1105           (
1106             'https://host1/test',
1107             {'data': data},
1108             json.dumps(gen_trigger_response()),
1109             None,
1110           ),
1111         ])
1112     ret = main([
1113         'trigger',
1114         '--swarming', 'https://host1',
1115         '--isolate-server', 'https://host2',
1116         '--shards', '1',
1117         '--priority', '101',
1118         '--dimension', 'foo', 'bar',
1119         '--dimension', 'os', 'Mac',
1120         '--task-name', TEST_NAME,
1121         '--deadline', '3600',
1122         FILE_HASH,
1123       ])
1124     actual = sys.stdout.getvalue()
1125     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
1126
1127   def test_trigger_dump_json(self):
1128     called = []
1129     self.mock(swarming.tools, 'write_json', lambda *args: called.append(args))
1130     self.mock(swarming.isolateserver, 'get_storage',
1131         lambda *_: MockedStorage(warm_cache=False))
1132     j = generate_expected_json(
1133         shards=1,
1134         shard_index=0,
1135         dimensions={'foo': 'bar', 'os': 'Mac'},
1136         env={},
1137         isolate_server='https://host2',
1138         profile=False)
1139     j['data'] = [[FAKE_BUNDLE_URL, 'swarm_data.zip']]
1140     data = {
1141       'request': json.dumps(j, sort_keys=True, separators=(',',':')),
1142     }
1143     self.expected_requests(
1144         [
1145           (
1146             'https://host1/test',
1147             {'data': data},
1148             json.dumps(gen_trigger_response()),
1149             None,
1150           ),
1151         ])
1152     ret = main([
1153         'trigger',
1154         '--swarming', 'https://host1',
1155         '--isolate-server', 'https://host2',
1156         '--shards', '1',
1157         '--priority', '101',
1158         '--dimension', 'foo', 'bar',
1159         '--dimension', 'os', 'Mac',
1160         '--task-name', TEST_NAME,
1161         '--deadline', '3600',
1162         '--dump-json', 'foo.json',
1163         FILE_HASH,
1164       ])
1165     actual = sys.stdout.getvalue()
1166     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
1167     expected = [
1168       (
1169         'foo.json',
1170         {
1171           u'base_task_name': u'unit_tests',
1172           u'tasks': {
1173             u'unit_tests': {
1174               u'shard_index': 0,
1175               u'task_id': u'123',
1176               u'view_url': u'https://host1/user/task/123',
1177             }
1178           },
1179         },
1180         True,
1181       ),
1182     ]
1183     self.assertEqual(expected, called)
1184
1185   def test_query_base(self):
1186     self.expected_requests(
1187         [
1188           (
1189             'https://localhost:1/swarming/api/v1/client/bots/botid/tasks?'
1190                 'limit=200',
1191             {},
1192             {'yo': 'dawg'},
1193           ),
1194         ])
1195     main(
1196         [
1197           'query', '--swarming', 'https://localhost:1', 'bots/botid/tasks',
1198         ])
1199     self._check_output('{\n  "yo": "dawg"\n}\n', '')
1200
1201   def test_query_cursor(self):
1202     self.expected_requests(
1203         [
1204           (
1205             'https://localhost:1/swarming/api/v1/client/bots/botid/tasks?'
1206                 'limit=2',
1207             {},
1208             {
1209               'cursor': '%',
1210               'extra': False,
1211               'items': ['A'],
1212             },
1213           ),
1214           (
1215             'https://localhost:1/swarming/api/v1/client/bots/botid/tasks?'
1216                 'cursor=%25&limit=1',
1217             {},
1218             {
1219               'cursor': None,
1220               'items': ['B'],
1221               'ignored': True,
1222             },
1223           ),
1224         ])
1225     main(
1226         [
1227           'query', '--swarming', 'https://localhost:1', 'bots/botid/tasks',
1228           '--limit', '2',
1229         ])
1230     expected = (
1231         '{\n'
1232         '  "extra": false, \n'
1233         '  "items": [\n'
1234         '    "A", \n'
1235         '    "B"\n'
1236         '  ]\n'
1237         '}\n')
1238     self._check_output(expected, '')
1239
1240
1241 class BotTestCase(TestCase):
1242   def setUp(self):
1243     super(BotTestCase, self).setUp()
1244     # Expected requests are always the same, independent of the test case.
1245     self.expected_requests(
1246         [
1247           (
1248             'https://localhost:1/swarming/api/v1/client/bots?limit=250',
1249             {},
1250             self.mock_swarming_api_v1_bots_page_1(),
1251           ),
1252           (
1253             'https://localhost:1/swarming/api/v1/client/bots?limit=250&'
1254               'cursor=opaque_cursor',
1255             {},
1256             self.mock_swarming_api_v1_bots_page_2(),
1257           ),
1258         ])
1259
1260   @staticmethod
1261   def mock_swarming_api_v1_bots_page_1():
1262     """Returns fake /swarming/api/v1/client/bots data."""
1263     # Sample data retrieved from actual server.
1264     now = unicode(datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S'))
1265     return {
1266       u'items': [
1267         {
1268           u'created_ts': now,
1269           u'dimensions': {
1270             u'cores': u'4',
1271             u'cpu': [u'x86', u'x86-64'],
1272             u'gpu': [u'15ad', u'15ad:0405'],
1273             u'hostname': u'swarm3.example.com',
1274             u'id': u'swarm3',
1275             u'os': [u'Mac', u'Mac-10.9'],
1276           },
1277           u'external_ip': u'1.1.1.3',
1278           u'hostname': u'swarm3.example.com',
1279           u'id': u'swarm3',
1280           u'internal_ip': u'192.168.0.3',
1281           u'is_dead': False,
1282           u'last_seen_ts': now,
1283           u'quarantined': False,
1284           u'task': u'148569b73a89501',
1285           u'version': u'56918a2ea28a6f51751ad14cc086f118b8727905',
1286         },
1287         {
1288           u'created_ts': now,
1289           u'dimensions': {
1290             u'cores': u'8',
1291             u'cpu': [u'x86', u'x86-64'],
1292             u'gpu': [],
1293             u'hostname': u'swarm1.example.com',
1294             u'id': u'swarm1',
1295             u'os': [u'Linux', u'Linux-12.04'],
1296           },
1297           u'external_ip': u'1.1.1.1',
1298           u'hostname': u'swarm1.example.com',
1299           u'id': u'swarm1',
1300           u'internal_ip': u'192.168.0.1',
1301           u'is_dead': True,
1302           u'last_seen_ts': 'A long time ago',
1303           u'quarantined': False,
1304           u'task': None,
1305           u'version': u'56918a2ea28a6f51751ad14cc086f118b8727905',
1306         },
1307         {
1308           u'created_ts': now,
1309           u'dimensions': {
1310             u'cores': u'8',
1311             u'cpu': [u'x86', u'x86-64'],
1312             u'cygwin': u'0',
1313             u'gpu': [
1314               u'15ad',
1315               u'15ad:0405',
1316               u'VMware Virtual SVGA 3D Graphics Adapter',
1317             ],
1318             u'hostname': u'swarm2.example.com',
1319             u'id': u'swarm2',
1320             u'integrity': u'high',
1321             u'os': [u'Windows', u'Windows-6.1'],
1322           },
1323           u'external_ip': u'1.1.1.2',
1324           u'hostname': u'swarm2.example.com',
1325           u'id': u'swarm2',
1326           u'internal_ip': u'192.168.0.2',
1327           u'is_dead': False,
1328           u'last_seen_ts': now,
1329           u'quarantined': False,
1330           u'task': None,
1331           u'version': u'56918a2ea28a6f51751ad14cc086f118b8727905',
1332         },
1333       ],
1334       u'cursor': u'opaque_cursor',
1335       u'death_timeout': 1800.0,
1336       u'limit': 4,
1337       u'now': unicode(now),
1338     }
1339
1340   @staticmethod
1341   def mock_swarming_api_v1_bots_page_2():
1342     """Returns fake /swarming/api/v1/client/bots data."""
1343     # Sample data retrieved from actual server.
1344     now = unicode(datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S'))
1345     return {
1346       u'items': [
1347         {
1348           u'created_ts': now,
1349           u'dimensions': {
1350             u'cores': u'8',
1351             u'cpu': [u'x86', u'x86-64'],
1352             u'gpu': [],
1353             u'hostname': u'swarm4.example.com',
1354             u'id': u'swarm4',
1355             u'os': [u'Linux', u'Linux-12.04'],
1356           },
1357           u'external_ip': u'1.1.1.4',
1358           u'hostname': u'swarm4.example.com',
1359           u'id': u'swarm4',
1360           u'internal_ip': u'192.168.0.4',
1361           u'is_dead': False,
1362           u'last_seen_ts': now,
1363           u'quarantined': False,
1364           u'task': u'14856971a64c601',
1365           u'version': u'56918a2ea28a6f51751ad14cc086f118b8727905',
1366         }
1367       ],
1368       u'cursor': None,
1369       u'death_timeout': 1800.0,
1370       u'limit': 4,
1371       u'now': unicode(now),
1372     }
1373
1374   def test_bots(self):
1375     main(['bots', '--swarming', 'https://localhost:1'])
1376     expected = (
1377         u'swarm2\n'
1378         u'  {"cores": "8", "cpu": ["x86", "x86-64"], "cygwin": "0", "gpu": '
1379           '["15ad", "15ad:0405", "VMware Virtual SVGA 3D Graphics Adapter"], '
1380           '"hostname": "swarm2.example.com", "id": "swarm2", "integrity": '
1381           '"high", "os": ["Windows", "Windows-6.1"]}\n'
1382         'swarm3\n'
1383         '  {"cores": "4", "cpu": ["x86", "x86-64"], "gpu": ["15ad", '
1384           '"15ad:0405"], "hostname": "swarm3.example.com", "id": "swarm3", '
1385           '"os": ["Mac", "Mac-10.9"]}\n'
1386         u'  task: 148569b73a89501\n'
1387         u'swarm4\n'
1388         u'  {"cores": "8", "cpu": ["x86", "x86-64"], "gpu": [], "hostname": '
1389           '"swarm4.example.com", "id": "swarm4", "os": ["Linux", '
1390           '"Linux-12.04"]}\n'
1391         u'  task: 14856971a64c601\n')
1392     self._check_output(expected, '')
1393
1394   def test_bots_bare(self):
1395     main(['bots', '--swarming', 'https://localhost:1', '--bare'])
1396     self._check_output("swarm2\nswarm3\nswarm4\n", '')
1397
1398   def test_bots_filter(self):
1399     main(
1400         [
1401           'bots', '--swarming', 'https://localhost:1',
1402           '--dimension', 'os', 'Windows',
1403         ])
1404     expected = (
1405         u'swarm2\n  {"cores": "8", "cpu": ["x86", "x86-64"], "cygwin": "0", '
1406           '"gpu": ["15ad", "15ad:0405", "VMware Virtual SVGA 3D Graphics '
1407           'Adapter"], "hostname": "swarm2.example.com", "id": "swarm2", '
1408           '"integrity": "high", "os": ["Windows", "Windows-6.1"]}\n')
1409     self._check_output(expected, '')
1410
1411   def test_bots_filter_keep_dead(self):
1412     main(
1413         [
1414           'bots', '--swarming', 'https://localhost:1',
1415           '--dimension', 'os', 'Linux', '--keep-dead',
1416         ])
1417     expected = (
1418         u'swarm1\n  {"cores": "8", "cpu": ["x86", "x86-64"], "gpu": [], '
1419           '"hostname": "swarm1.example.com", "id": "swarm1", "os": ["Linux", '
1420           '"Linux-12.04"]}\n'
1421         u'swarm4\n'
1422         u'  {"cores": "8", "cpu": ["x86", "x86-64"], "gpu": [], "hostname": '
1423           '"swarm4.example.com", "id": "swarm4", "os": ["Linux", '
1424           '"Linux-12.04"]}\n'
1425         u'  task: 14856971a64c601\n')
1426     self._check_output(expected, '')
1427
1428   def test_bots_filter_dead_only(self):
1429     main(
1430         [
1431           'bots', '--swarming', 'https://localhost:1',
1432           '--dimension', 'os', 'Linux', '--dead-only',
1433         ])
1434     expected = (
1435         u'swarm1\n  {"cores": "8", "cpu": ["x86", "x86-64"], "gpu": [], '
1436           '"hostname": "swarm1.example.com", "id": "swarm1", "os": ["Linux", '
1437           '"Linux-12.04"]}\n')
1438     self._check_output(expected, '')
1439
1440
1441 def gen_run_isolated_out_hack_log(isolate_server, namespace, isolated_hash):
1442   data = {
1443     'hash': isolated_hash,
1444     'namespace': namespace,
1445     'storage': isolate_server,
1446   }
1447   return (SWARM_OUTPUT_SUCCESS +
1448       '[run_isolated_out_hack]%s[/run_isolated_out_hack]\n' % (
1449           json.dumps(data, sort_keys=True, separators=(',',':'))))
1450
1451
1452 class ExtractOutputFilesLocationTest(auto_stub.TestCase):
1453   def test_ok(self):
1454     task_log = '\n'.join((
1455       'some log',
1456       'some more log',
1457       gen_run_isolated_out_hack_log('https://fake', 'default', '12345'),
1458       'more log',
1459     ))
1460     self.assertEqual(
1461         {'hash': '12345',
1462          'namespace': 'default',
1463          'server': 'https://fake',
1464          'view_url': 'https://fake/browse?namespace=default&hash=12345'},
1465         swarming.extract_output_files_location(task_log))
1466
1467   def test_empty(self):
1468     task_log = '\n'.join((
1469       'some log',
1470       'some more log',
1471       '[run_isolated_out_hack]',
1472       '[/run_isolated_out_hack]',
1473     ))
1474     self.assertEqual(
1475         None,
1476         swarming.extract_output_files_location(task_log))
1477
1478   def test_missing(self):
1479     task_log = '\n'.join((
1480       'some log',
1481       'some more log',
1482       'more log',
1483     ))
1484     self.assertEqual(
1485         None,
1486         swarming.extract_output_files_location(task_log))
1487
1488   def test_corrupt(self):
1489     task_log = '\n'.join((
1490       'some log',
1491       'some more log',
1492       '[run_isolated_out_hack]',
1493       '{"hash": "12345","namespace":}',
1494       '[/run_isolated_out_hack]',
1495       'more log',
1496     ))
1497     self.assertEqual(
1498         None,
1499         swarming.extract_output_files_location(task_log))
1500
1501   def test_not_url(self):
1502     task_log = '\n'.join((
1503       'some log',
1504       'some more log',
1505       gen_run_isolated_out_hack_log('/local/path', 'default', '12345'),
1506       'more log',
1507     ))
1508     self.assertEqual(
1509         None,
1510         swarming.extract_output_files_location(task_log))
1511
1512
1513 class TaskOutputCollectorTest(auto_stub.TestCase):
1514   def setUp(self):
1515     super(TaskOutputCollectorTest, self).setUp()
1516
1517     # Silence error log.
1518     self.mock(logging, 'error', lambda *_, **__: None)
1519
1520     # Collect calls to 'isolateserver.fetch_isolated'.
1521     self.fetch_isolated_calls = []
1522     def fetch_isolated(isolated_hash, storage, cache, outdir, require_command):
1523       self.fetch_isolated_calls.append(
1524           (isolated_hash, storage, cache, outdir, require_command))
1525     # Ensure mock has exact same signature as the original, otherwise tests may
1526     # miss changes to real 'fetch_isolated' arg list.
1527     self.assertEqual(
1528         inspect.getargspec(swarming.isolateserver.fetch_isolated),
1529         inspect.getargspec(fetch_isolated))
1530     self.mock(swarming.isolateserver, 'fetch_isolated', fetch_isolated)
1531
1532     # TaskOutputCollector creates directories. Put them in a temp directory.
1533     self.tempdir = tempfile.mkdtemp(prefix='swarming_test')
1534
1535   def tearDown(self):
1536     shutil.rmtree(self.tempdir)
1537     super(TaskOutputCollectorTest, self).tearDown()
1538
1539   def test_works(self):
1540     # Output logs of shards.
1541     logs = [
1542       gen_run_isolated_out_hack_log('https://server', 'namespace', 'hash1'),
1543       gen_run_isolated_out_hack_log('https://server', 'namespace', 'hash2'),
1544       SWARM_OUTPUT_SUCCESS,
1545     ]
1546
1547     # Feed three shard results to collector, last one without output files.
1548     collector = swarming.TaskOutputCollector(
1549         self.tempdir, 'task/name', len(logs))
1550     for index, log in enumerate(logs):
1551       collector.process_shard_result(index, gen_data(log, '0, 0'))
1552     summary = collector.finalize()
1553
1554     # Ensure it fetches the files from first two shards only.
1555     expected_calls = [
1556       ('hash1', None, None, os.path.join(self.tempdir, '0'), False),
1557       ('hash2', None, None, os.path.join(self.tempdir, '1'), False),
1558     ]
1559     self.assertEqual(len(expected_calls), len(self.fetch_isolated_calls))
1560     storage_instances = set()
1561     for expected, used in zip(expected_calls, self.fetch_isolated_calls):
1562       isolated_hash, storage, cache, outdir, require_command = used
1563       storage_instances.add(storage)
1564       # Compare everything but |storage| and |cache| (use None in their place).
1565       self.assertEqual(
1566           expected, (isolated_hash, None, None, outdir, require_command))
1567       # Ensure cache is set.
1568       self.assertTrue(cache)
1569
1570     # Only one instance of Storage should be used.
1571     self.assertEqual(1, len(storage_instances))
1572
1573     # Ensure storage is pointing to required location.
1574     storage = storage_instances.pop()
1575     self.assertEqual('https://server', storage.location)
1576     self.assertEqual('namespace', storage.namespace)
1577
1578     # Ensure collected summary is correct.
1579     expected_summary = {
1580       'task_name': 'task/name',
1581       'shards': [
1582         gen_data(log, '0, 0') for index, log in enumerate(logs)
1583       ]
1584     }
1585     self.assertEqual(expected_summary, summary)
1586
1587     # Ensure summary dumped to a file is correct as well.
1588     with open(os.path.join(self.tempdir, 'summary.json'), 'r') as f:
1589       summary_dump = json.load(f)
1590     self.assertEqual(expected_summary, summary_dump)
1591
1592   def test_ensures_same_server(self):
1593     # Two shard results, attempt to use different servers.
1594     data = [
1595       gen_data(
1596         gen_run_isolated_out_hack_log('https://server1', 'namespace', 'hash1'),
1597         '0, 0'),
1598       gen_data(
1599         gen_run_isolated_out_hack_log('https://server2', 'namespace', 'hash2'),
1600         '0, 0'),
1601     ]
1602
1603     # Feed them to collector.
1604     collector = swarming.TaskOutputCollector(self.tempdir, 'task/name', 2)
1605     for index, result in enumerate(data):
1606       collector.process_shard_result(index, result)
1607     collector.finalize()
1608
1609     # Only first fetch is made, second one is ignored.
1610     self.assertEqual(1, len(self.fetch_isolated_calls))
1611     isolated_hash, storage, _, outdir, _ = self.fetch_isolated_calls[0]
1612     self.assertEqual(
1613         ('hash1', os.path.join(self.tempdir, '0')),
1614         (isolated_hash, outdir))
1615     self.assertEqual('https://server1', storage.location)
1616
1617
1618 def clear_env_vars():
1619   for e in ('ISOLATE_SERVER', 'SWARMING_SERVER'):
1620     os.environ.pop(e, None)
1621
1622
1623 if __name__ == '__main__':
1624   logging.basicConfig(
1625       level=logging.DEBUG if '-v' in sys.argv else logging.ERROR)
1626   if '-v' in sys.argv:
1627     unittest.TestCase.maxDiff = None
1628   clear_env_vars()
1629   unittest.main()