Reapply "Isolated task support in Endpoints API: client side (3/3)" and fixes"

appengine/swarming/local_smoke_test.py

 #!/usr/bin/env python
 # Copyright 2014 The Swarming Authors. All rights reserved.
 # Use of this source code is governed by the Apache v2.0 license 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 32 matching lines @@
 SIGNAL_TERM = -1073741510 if sys.platform == 'win32' else -signal.SIGTERM
 
 
 class SwarmingClient(object):
   def __init__(self, swarming_server, isolate_server):
     self._swarming_server = swarming_server
     self._isolate_server = isolate_server
     self._tmpdir = tempfile.mkdtemp(prefix='swarming_client')
     self._index = 0
 
+  def isolate(self, isolate_path, isolated_path):
+    cmd = [
+      sys.executable, 'isolate.py', 'archive',
+      '-I', self._isolate_server,
+      '--namespace', 'default-gzip',
+      '-i', isolate_path,
+      '-s', isolated_path,
+    ]
+    isolated_hash = subprocess.check_output(cmd, cwd=CLIENT_DIR).split()[0]
+    logging.debug('%s = %s', isolated_path, isolated_hash)
+    return isolated_hash
+
   def task_trigger_raw(self, args):
     """Triggers a task and return the task id."""
     h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
     os.close(h)
     try:
       cmd = [
         '--user', 'joe@localhost',
         '-d', 'cpu', 'x86',
         '--dump-json', tmp,
         '--raw-cmd',
       ]
       cmd.extend(args)
       assert not self._run('trigger', cmd), args
       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:
       os.remove(tmp)
 
+  def task_trigger_isolated(self, name, isolated_hash, extra=None):
+    """Triggers a task and return the task id."""
+    h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
+    os.close(h)
+    try:
+      cmd = [
+        '--user', 'joe@localhost',
+        '-d', 'cpu', 'x86',
+        '--dump-json', tmp,
+        '--task-name', name,
+        '-I',  self._isolate_server,
+        '--namespace', 'default-gzip',
+        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:
+      os.remove(tmp)
+
   def task_collect(self, task_id):
     """Collects the results for a task."""
     h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
     os.close(h)
     try:
       tmpdir = tempfile.mkdtemp(prefix='swarming_smoke_test')
       try:
         # swarming.py collect will return the exit code of the task.
         args = [
           '--task-summary-json', tmp, task_id, '--task-output-dir', tmpdir,
@@ skipping 207 matching lines @@
     # This will restart the bot. This ensures the update mechanism works.
     # TODO(maruel): Convert to a real API. Can only be accessed by admin-level
     # account.
     res = self.servers.http_client.request(
         '/restricted/upload/bot_config',
         body=urllib.urlencode({'script': bot_config_content}))
     self.assertEqual(200, res.http_code, res.body)
     bot_version2 = self.assertOneTask(args, summary, {})
     self.assertNotEqual(bot_version1, bot_version2)
 
+  def test_isolated(self):
+    # Make an isolated file, archive it.
+    isolate = {
+      'variables': {
+        'command': ['python', 'hello_world.py'],
+        'files': ['hello_world.py'],
+      },
+    }
+    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(\'hey\')'))
+    expected_summary = self.gen_expected(
+        name=u'yo',
+        isolated_out={
+          u'isolated': u'f10f4c42b38ca01726610f9575ba695468c32108',
+          u'isolatedserver': u'http://localhost:10050',
+          u'namespace': u'default-gzip',
+          u'view_url':
+            u'http://localhost:10050/browse?namespace=default-gzip'
+            '&hash=f10f4c42b38ca01726610f9575ba695468c32108',
+        },
+        outputs=[
+          u'hi\n'
+        ],
+        outputs_ref={
+          u'isolated': u'f10f4c42b38ca01726610f9575ba695468c32108',
+          u'isolatedserver': u'http://localhost:10050',
+          u'namespace': u'default-gzip',
+          u'view_url':
+            u'http://localhost:10050/browse?namespace=default-gzip'
+            '&hash=f10f4c42b38ca01726610f9575ba695468c32108',
+        })
+    expected_files = {os.path.join('0', 'result.txt'): 'hey'}
+    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, 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(
+          'yo', isolated_hash, extra=['--', '${ISOLATED_OUTDIR}'])
+      actual_summary, actual_files = self.client.task_collect(task_id)
+      self.assertResults(expected_summary, actual_summary)
+      actual_files.pop('summary.json')
+      self.assertEqual(expected_files, actual_files)
+    finally:
+      shutil.rmtree(tmpdir)
+
   def assertResults(self, expected, result):
     self.assertEqual(['shards'], result.keys())
     self.assertEqual(1, len(result['shards']))
     self.assertTrue(result['shards'][0])
     result = result['shards'][0].copy()
     # These are not deterministic (or I'm too lazy to calculate the value).
     bot_version = result.pop('bot_version')
     self.assertTrue(bot_version)
     self.assertTrue(result.pop('costs_usd'))
     self.assertTrue(result.pop('created_ts'))
@@ skipping 71 matching lines @@
     print >> sys.stderr, '<Ctrl-C>'
     if bot:
       bot.kill()
   finally:
     cleanup(bot, client, servers, failed or verbose)
   return int(failed)
 
 
 if __name__ == '__main__':
   sys.exit(main())