swarming: fix support for raw command + isolated file

appengine/swarming/local_smoke_test.py

 #!/usr/bin/env python
 # Copyright 2014 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.
 
 """Integration test for the Swarming server, Swarming bot and Swarming client.
 
 It starts both a Swarming server and a Swarming bot and triggers tasks with the
 Swarming client to ensure the system works end to end.
 """
@@ skipping 56 matching lines @@
 
 RESULT_HEY_OUTPUTS_REF = {
   u'isolated': u'5c64883277eb00bceafe3659f182e194cffc5d96',
   u'isolatedserver': u'http://localhost:10050',
   u'namespace': u'default-gzip',
   u'view_url':
     u'http://localhost:10050/browse?namespace=default-gzip'
     '&hash=5c64883277eb00bceafe3659f182e194cffc5d96',
 }
 
+RESULT_HEY2_ISOLATED_OUT = {
+  u'isolated': u'0616f86b24065a0595e58088925567ae54a8157c',
+  u'isolatedserver': u'http://localhost:10050',
+  u'namespace': u'default-gzip',
+  u'view_url':
+    u'http://localhost:10050/browse?namespace=default-gzip'
+    '&hash=0616f86b24065a0595e58088925567ae54a8157c',
+}
+
+RESULT_HEY2_OUTPUTS_REF = {
+  u'isolated': u'0616f86b24065a0595e58088925567ae54a8157c',
+  u'isolatedserver': u'http://localhost:10050',
+  u'namespace': u'default-gzip',
+  u'view_url':
+    u'http://localhost:10050/browse?namespace=default-gzip'
+    '&hash=0616f86b24065a0595e58088925567ae54a8157c',
+}
+
 RESULT_SECRET_OUTPUT = {
   u'isolated': u'd2eca4d860e4f1728272f6a736fd1c9ac6e98c4f',
   u'isolatedserver': u'http://localhost:10050',
   u'namespace': u'default-gzip',
   u'view_url':
     u'http://localhost:10050/browse?namespace=default-gzip'
     '&hash=d2eca4d860e4f1728272f6a736fd1c9ac6e98c4f',
 }
 
 
@@ skipping 42 matching lines @@
     h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
     os.close(h)
     try:
       cmd = [
         '--user', 'joe@localhost',
         '-d', 'pool', 'default',
         '--dump-json', tmp,
         '--task-name', name,
         '-I',  self._isolate_server,
         '--namespace', 'default-gzip',
-        isolated_hash,
+        '-s', isolated_hash,
       ]
       if extra:
         cmd.extend(extra)
       assert not self._run('trigger', cmd)
       with open(tmp, 'rb') as f:
         data = json.load(f)
         task_id = data['tasks'].popitem()[1]['task_id']
         logging.debug('task_id = %s', task_id)
         return task_id
     finally:
@@ skipping 117 matching lines @@
   bot = None
 
   @classmethod
   def setUpClass(cls):
     cls.dimensions = os_utilities.get_dimensions()
 
   def setUp(self):
     super(Test, self).setUp()
     # Reset the bot's cache at the start of each task, so that the cache reuse
     # data becomes deterministic.
+    # Main caveat is 'isolated_upload' as the isolate server is not cleared.
     self.bot.wipe_cache()
 
   def gen_expected(self, **kwargs):
     return gen_expected(bot_dimensions=self.dimensions, **kwargs)
 
   def test_raw_bytes(self):
     # A string of a letter 'A', UTF-8 BOM then UTF-16 BOM then UTF-EDBCDIC then
     # invalid UTF-8 and the letter 'B'. It is double escaped so it can be passed
     # down the shell.
     invalid_bytes = 'A\\xEF\\xBB\\xBF\\xFE\\xFF\\xDD\\x73\\x66\\x73\\xc3\\x28B'
@@ skipping 114 matching lines @@
         'print(\'hi\')',
         'with open(os.path.join(sys.argv[1], \'result.txt\'), \'wb\') as f:',
         '  f.write(\'hey\')'))
     expected_summary = self.gen_expected(
         name=u'isolated_task',
         isolated_out=RESULT_HEY_ISOLATED_OUT,
         performance_stats={
           u'isolated_download': {
             u'initial_number_items': u'0',
             u'initial_size': u'0',
-            u'items_cold': [112, 200],
+            u'items_cold': sorted([len(hello_world), 200]),
             u'items_hot': [],
           },
           u'isolated_upload': {
             u'items_cold': [3, 118],
             u'items_hot': [],
           },
         },
         outputs=[u'hi\n'],
         outputs_ref=RESULT_HEY_OUTPUTS_REF)
     expected_files = {os.path.join('0', 'result.txt'): 'hey'}
     self._run_isolated(
         hello_world, 'isolated_task', ['--', '${ISOLATED_OUTDIR}'],
         expected_summary, expected_files)
 
+  def test_isolated_command(self):
+    # Command is specified in Swarming task, still with isolated file.
+    hello_world = '\n'.join((
+        'import os',
+        'import sys',
+        'print(\'hi\')',
+        'with open(os.path.join(sys.argv[1], \'result.txt\'), \'wb\') as f:',
+        '  f.write(\'hey2\')'))
+    expected_summary = self.gen_expected(
+        name=u'separate_cmd',
+        isolated_out=RESULT_HEY2_ISOLATED_OUT,
+        performance_stats={
+          u'isolated_download': {
+            u'initial_number_items': u'0',
+            u'initial_size': u'0',
+            u'items_cold': sorted([len(hello_world), 157]),
+            u'items_hot': [],
+          },
+          u'isolated_upload': {
+            u'items_cold': [4, 118],
+            u'items_hot': [],
+          },
+        },
+        outputs=[u'hi\n'],
+        outputs_ref=RESULT_HEY2_OUTPUTS_REF)
+    expected_files = {os.path.join('0', 'result.txt'): 'hey2'}
+    self._run_isolated(
+        hello_world, 'separate_cmd',
+        ['--raw-cmd', '--', 'python', 'hello_world.py', '${ISOLATED_OUTDIR}'],
+        expected_summary, expected_files,
+        isolated_content={'variables': {'files': ['hello_world.py']}})
+
   def test_isolated_hard_timeout(self):
     # Make an isolated file, archive it, have it time out. Similar to
     # test_hard_timeout. The script doesn't handle signal so it failed the grace
     # period.
     hello_world = '\n'.join((
         'import os',
         'import sys',
         'import time',
         'sys.stdout.write(\'hi\\n\')',
         'sys.stdout.flush()',
         'time.sleep(120)',
         'with open(os.path.join(sys.argv[1], \'result.txt\'), \'wb\') as f:',
         '  f.write(\'hey\')'))
     expected_summary = self.gen_expected(
         name=u'isolated_hard_timeout',
         exit_codes=[SIGNAL_TERM],
         failure=True,
         performance_stats={
           u'isolated_download': {
             u'initial_number_items': u'0',
             u'initial_size': u'0',
-            u'items_cold': [172, 200],
+            u'items_cold': sorted([len(hello_world), 200]),
             u'items_hot': [],
           },
           u'isolated_upload': {
             u'items_cold': [],
             u'items_hot': [],
           },
         },
         state=0x40)  # task_result.State.TIMED_OUT
     # Hard timeout is enforced by run_isolated, I/O timeout by task_runner.
     self._run_isolated(
@@ skipping 25 matching lines @@
         '    print(\'ioerror\')',
         'with open(os.path.join(sys.argv[1], \'result.txt\'), \'wb\') as f:',
         '  f.write(\'hey\')'))
     expected_summary = self.gen_expected(
         name=u'isolated_hard_timeout_grace',
         isolated_out=RESULT_HEY_ISOLATED_OUT,
         performance_stats={
           u'isolated_download': {
             u'initial_number_items': u'0',
             u'initial_size': u'0',
-            u'items_cold': [200, 407],
+            u'items_cold': sorted([200, len(hello_world)]),
             u'items_hot': [],
           },
           u'isolated_upload': {
             u'items_cold': [],
             u'items_hot': [3, 118],
           },
         },
         outputs=[u'hi\ngot signal 15\n'],
         outputs_ref=RESULT_HEY_OUTPUTS_REF,
         failure=True,
         state=0x40)  # task_result.State.TIMED_OUT
     expected_files = {os.path.join('0', 'result.txt'): 'hey'}
     # Hard timeout is enforced by run_isolated, I/O timeout by task_runner.
     self._run_isolated(
         hello_world, 'isolated_hard_timeout_grace',
         ['--hard-timeout', '1', '--', '${ISOLATED_OUTDIR}'],
         expected_summary, expected_files)
 
   def test_idempotent_reuse(self):
     hello_world = 'print "hi"\n'
     expected_summary = self.gen_expected(
       name=u'idempotent_reuse',
       performance_stats={
         u'isolated_download': {
           u'initial_number_items': u'0',
           u'initial_size': u'0',
-          u'items_cold': [11, 199],
+          u'items_cold': sorted([len(hello_world), 199]),
           u'items_hot': [],
         },
         u'isolated_upload': {
           u'items_cold': [],
           u'items_hot': [],
         },
       },
       properties_hash =
           u'082928de84d0a65839d227dcea2f5a947898929c77c1602b68c46d7d4588c1f5',
     )
@@ skipping 23 matching lines @@
       with open(os.path.join(sys.argv[1], 'sekret'), 'w') as f:
         print >> f, data['swarming']['secret_bytes'].decode('base64')
     """)
     expected_summary = self.gen_expected(
       name=u'secret_bytes',
       isolated_out=RESULT_SECRET_OUTPUT,
       performance_stats={
         u'isolated_download': {
           u'initial_number_items': u'0',
           u'initial_size': u'0',
-          u'items_cold': [200, 241],
+          u'items_cold': sorted([200, len(hello_world)]),
           u'items_hot': [],
         },
         u'isolated_upload': {
           u'items_cold': [7, 114],
           u'items_hot': [],
         },
       },
       outputs_ref=RESULT_SECRET_OUTPUT,
     )
     h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test_secret')
@@ skipping 69 matching lines @@
         },
       },
     )
     self._run_isolated(
         script, 'cache_second',
         ['--named-cache', 'fuu', 'p/b', '--', '${ISOLATED_OUTDIR}/yo'],
         expected_summary,
         {'0/yo': 'Yo!'})
 
   def _run_isolated(self, hello_world, name, args, expected_summary,
-      expected_files, deduped=False):
+      expected_files, deduped=False, isolated_content=None):
+    """Runs hello_world.py as an isolated file."""
     # Shared code for all test_isolated_* test cases.
     tmpdir = tempfile.mkdtemp(prefix='swarming_smoke')
     try:
       isolate_path = os.path.join(tmpdir, 'i.isolate')
       isolated_path = os.path.join(tmpdir, 'i.isolated')
       with open(isolate_path, 'wb') as f:
-        json.dump(ISOLATE_HELLO_WORLD, f)
+        json.dump(isolated_content or ISOLATE_HELLO_WORLD, f)
       with open(os.path.join(tmpdir, 'hello_world.py'), 'wb') as f:
         f.write(hello_world)
       isolated_hash = self.client.isolate(isolate_path, isolated_path)
       task_id = self.client.task_trigger_isolated(
           name, isolated_hash, extra=args)
       actual_summary, actual_files = self.client.task_collect(task_id)
       self.assertResults(expected_summary, actual_summary, deduped=deduped)
       actual_files.pop('summary.json')
       self.assertEqual(expected_files, actual_files)
       return task_id
@@ skipping 149 matching lines @@
     if bot is not None and bot.poll() is None:
       bot.kill()
       bot.wait()
   finally:
     cleanup(bot, client, servers, failed or verbose, leak)
   return int(failed)
 
 
 if __name__ == '__main__':
   sys.exit(main())