swarming: fix support for raw command + isolated file

client/tests/swarming_test.py

 #!/usr/bin/env python
 # Copyright 2013 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 import datetime
 import json
 import logging
 import os
 import re
 import StringIO
 import subprocess
 import sys
 import tempfile
 import threading
 import time
+import traceback
 import unittest
 
 # net_utils adjusts sys.path.
 import net_utils
 
 from depot_tools import auto_stub
 
 import auth
 import isolateserver
 import swarming
@@ skipping 68 matching lines @@
       'command': None,
       'dimensions': [
         {'key': 'foo', 'value': 'bar'},
         {'key': 'os', 'value': 'Mac'},
       ],
       'env': [],
       'execution_timeout_secs': 60,
       'extra_args': ['--some-arg', '123'],
       'grace_period_secs': 30,
       'idempotent': False,
-      'inputs_ref': None,
+      'inputs_ref': {
+        'isolated': None,
+        'isolatedserver': '',
+        'namespace': 'default-gzip',
+      },
       'io_timeout_secs': 60,
       'outputs': [],
       'secret_bytes': None,
     },
     'tags': ['tag:a', 'tag:b'],
     'user': 'joe@localhost',
   }
   out.update(kwargs)
   out['properties'].update(properties or {})
   return out
@@ skipping 113 matching lines @@
 
   def main_safe(self, args):
     """Bypasses swarming.main()'s exception handling.
 
     It gets in the way when debugging test failures.
     """
     # pylint: disable=bare-except
     try:
       return main(args)
     except:
-      logging.exception('Unexpected exception thrown in main')
-      logging.error(
-          'STDOUT:\n%s\nSTDERR:\n%s',
-          sys.stdout.getvalue(), sys.stderr.getvalue())
-      self.fail()
+      data = '%s\nSTDOUT:\n%s\nSTDERR:\n%s' % (
+          traceback.format_exc(), sys.stdout.getvalue(), sys.stderr.getvalue())
+      self.fail(data)
 
 
 class NetTestCase(net_utils.TestCase, Common):
   """Base class that defines the url_open mock."""
   def setUp(self):
     net_utils.TestCase.setUp(self)
     Common.setUp(self)
     self.mock(time, 'sleep', lambda _: None)
     self.mock(subprocess, 'call', lambda *_: self.fail())
     self.mock(threading, 'Event', NonBlockingEvent)
@@ skipping 79 matching lines @@
         properties=swarming.TaskProperties(
             caches=[],
             cipd_input=None,
             command=['a', 'b'],
             dimensions={'foo': 'bar', 'os': 'Mac'},
             env={},
             execution_timeout_secs=60,
             extra_args=[],
             grace_period_secs=30,
             idempotent=False,
-            inputs_ref=None,
+            inputs_ref={
+              'isolated': None,
+              'isolatedserver': '',
+              'namespace': 'default-gzip',
+            },
             io_timeout_secs=60,
             outputs=[],
             secret_bytes=None),
         service_account_token=None,
         tags=['tag:a', 'tag:b'],
         user='joe@localhost')
 
     request_1 = swarming.task_request_to_raw_request(task_request, False)
     request_1['name'] = u'unit_tests:0:2'
     request_1['properties']['env'] = [
@@ skipping 50 matching lines @@
         properties=swarming.TaskProperties(
             caches=[],
             cipd_input=None,
             command=['a', 'b'],
             dimensions={'foo': 'bar', 'os': 'Mac'},
             env={},
             execution_timeout_secs=60,
             extra_args=[],
             grace_period_secs=30,
             idempotent=False,
-            inputs_ref=None,
+            inputs_ref={
+              'isolated': None,
+              'isolatedserver': '',
+              'namespace': 'default-gzip',
+            },
             io_timeout_secs=60,
             outputs=[],
             secret_bytes=None),
         service_account_token=None,
         tags=['tag:a', 'tag:b'],
         user='joe@localhost')
 
     request = swarming.task_request_to_raw_request(task_request, False)
     self.assertEqual('123', request['parent_task_id'])
 
@@ skipping 42 matching lines @@
                         path='path/to/package',
                         version='abc123')],
                 server=None),
             command=['a', 'b'],
             dimensions={'foo': 'bar', 'os': 'Mac'},
             env={},
             execution_timeout_secs=60,
             extra_args=[],
             grace_period_secs=30,
             idempotent=False,
-            inputs_ref=None,
+            inputs_ref={
+              'isolated': None,
+              'isolatedserver': '',
+              'namespace': 'default-gzip',
+            },
             io_timeout_secs=60,
             outputs=[],
             secret_bytes=None),
         service_account_token=None,
         tags=['tag:a', 'tag:b'],
         user='joe@localhost')
 
     request = swarming.task_request_to_raw_request(task_request, False)
     expected = {
       'client_package': None,
@@ skipping 418 matching lines @@
     actual = sys.stdout.getvalue()
     self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
     self._check_output(
         'Triggered task: None/foo=bar\n'
         'To collect results, use:\n'
         '  swarming.py collect -S https://localhost:1 12300\n'
         'Or visit:\n'
         '  https://localhost:1/user/task/12300\n',
         '')
 
+  def test_run_raw_cmd_isolated(self):
+    # Minimalist use.
+    request = {
+      'expiration_secs': 21600,
+      'name': u'None/foo=bar/' + FILE_HASH,
+      'parent_task_id': '',
+      'priority': 100,
+      'properties': {
+        'caches': [],
+        'cipd_input': None,
+        'command': ['python', '-c', 'print(\'hi\')'],
+        'dimensions': [
+          {'key': 'foo', 'value': 'bar'},
+        ],
+        'env': [],
+        'execution_timeout_secs': 3600,
+        'extra_args': None,
+        'grace_period_secs': 30,
+        'idempotent': False,
+        'inputs_ref': {
+          'isolated': FILE_HASH,
+          'isolatedserver': 'https://localhost:2',
+          'namespace': 'default-gzip',
+        },
+        'io_timeout_secs': 1200,
+        'outputs': [],
+        'secret_bytes': None,
+      },
+      'tags': [],
+      'user': None,
+    }
+    result = gen_request_response(request)
+    self.expected_requests(
+        [
+          (
+            'https://localhost:1/api/swarming/v1/tasks/new',
+            {'data': request},
+            result,
+          ),
+        ])
+    ret = self.main_safe([
+        'trigger',
+        '--swarming', 'https://localhost:1',
+        '--dimension', 'foo', 'bar',
+        '--raw-cmd',
+        '--isolate-server', 'https://localhost:2',
+        '--isolated', FILE_HASH,
+        '--',
+        'python',
+        '-c',
+        'print(\'hi\')',
+      ])
+    actual = sys.stdout.getvalue()
+    self.assertEqual(0, ret, (actual, sys.stderr.getvalue()))
+    self._check_output(
+        u'Triggered task: None/foo=bar/' + FILE_HASH + u'\n'
+        u'To collect results, use:\n'
+        u'  swarming.py collect -S https://localhost:1 12300\n'
+        u'Or visit:\n'
+        u'  https://localhost:1/user/task/12300\n',
+        u'')
+
   def test_run_raw_cmd_with_service_account(self):
     # Minimalist use.
     request = {
       'expiration_secs': 21600,
       'name': u'None/foo=bar',
       'parent_task_id': '',
       'priority': 100,
       'properties': {
         'caches': [],
         'cipd_input': None,
@@ skipping 46 matching lines @@
         '')
 
   def test_run_isolated_hash(self):
     # pylint: disable=unused-argument
     self.mock(swarming, 'now', lambda: 123456)
 
     request = gen_request_data(
         properties={
           'command': None,
           'inputs_ref': {
-            'isolated': u'1111111111111111111111111111111111111111',
+            'isolated': FILE_HASH,
             'isolatedserver': 'https://localhost:2',
             'namespace': 'default-gzip',
           },
           'secret_bytes': None,
         })
     result = gen_request_response(request)
     self.expected_requests(
         [
           (
             'https://localhost:1/api/swarming/v1/tasks/new',
@@ skipping 154 matching lines @@
             'client_package': None,
             'packages': [{
               'package_name': 'super/awesome/pkg',
               'path': 'path/to/pkg',
               'version': 'version:42',
             }],
             'server': None,
           },
           'command': None,
           'inputs_ref': {
-            'isolated': u'1111111111111111111111111111111111111111',
+            'isolated': FILE_HASH,
             'isolatedserver': 'https://localhost:2',
             'namespace': 'default-gzip',
           },
           'secret_bytes': None,
         })
     result = gen_request_response(request)
     self.expected_requests(
         [
           (
             'https://localhost:1/api/swarming/v1/tasks/new',
@@ skipping 37 matching lines @@
       main([
             'trigger', '--swarming', 'https://host',
             '--isolate-server', 'https://host', '-T', 'foo',
             '-d', 'os', 'amgia',
           ])
     self._check_output(
         '',
         'Usage: swarming.py trigger [options] (hash|isolated) '
           '[-- extra_args|raw command]\n'
         '\n'
-        'swarming.py: error: Use --isolated, --raw-cmd or \'--\' to pass '
-          'arguments to the called process.\n')
+        'swarming.py: error: Specify at least one of --raw-cmd or --isolated '
+        'or both\n')
 
   def test_trigger_no_env_vars(self):
     with self.assertRaises(SystemExit):
       main(['trigger'])
     self._check_output(
         '',
         'Usage: swarming.py trigger [options] (hash|isolated) '
           '[-- extra_args|raw command]'
         '\n\n'
         'swarming.py: error: --swarming is required.'
@@ skipping 13 matching lines @@
 
   def test_trigger_no_isolate_server(self):
     with self.assertRaises(SystemExit):
       with test_utils.EnvVars({'SWARMING_SERVER': 'https://host'}):
         main(['trigger', 'foo.isolated', '-d', 'os', 'amiga'])
     self._check_output(
         '',
         'Usage: swarming.py trigger [options] (hash|isolated) '
           '[-- extra_args|raw command]'
         '\n\n'
-        'swarming.py: error: --isolate-server is required.'
-        '\n')
+        'swarming.py: error: Specify at least one of --raw-cmd or --isolated '
+          'or both\n')
 
   def test_trigger_no_dimension(self):
     with self.assertRaises(SystemExit):
       main([
             'trigger', '--swarming', 'https://host', '--raw-cmd', '--', 'foo',
           ])
     self._check_output(
         '',
         'Usage: swarming.py trigger [options] (hash|isolated) '
           '[-- extra_args|raw command]'
@@ skipping 21 matching lines @@
           'dimensions': [
             {'key': 'foo', 'value': 'bar'},
             {'key': 'os', 'value': 'Mac'},
           ],
           'env': [],
           'execution_timeout_secs': 60,
           'extra_args': ['--some-arg', '123'],
           'grace_period_secs': 30,
           'idempotent': True,
           'inputs_ref': {
-            'isolated': '1'*40,
+            'isolated': FILE_HASH,
             'isolatedserver': 'https://localhost:2',
             'namespace': 'default-gzip',
             },
           'io_timeout_secs': 60,
           'secret_bytes': None,
         },
         'tags': ['tag:a', 'tag:b'],
         'user': 'joe@localhost',
       },
     }
@@ skipping 322 matching lines @@
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   logging.basicConfig(
       level=logging.DEBUG if '-v' in sys.argv else logging.CRITICAL)
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
   for e in ('ISOLATE_SERVER', 'SWARMING_TASK_ID', 'SWARMING_SERVER'):
     os.environ.pop(e, None)
   unittest.main()