Send authorization headers when calling Swarming backend.

appengine/swarming/swarming_bot/bot_code/task_runner_test.py

 #!/usr/bin/env python
 # coding=utf-8
 # 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 base64
 import json
 import logging
 import os
 import re
 import signal
 import sys
 import tempfile
 import time
 import unittest
 
 import test_env_bot_code
 test_env_bot_code.setup_test_env()
 
 # Creates a server mock for functions in net.py.
 import net_utils
 
-from api import os_utilities
 from depot_tools import fix_encoding
 from utils import file_path
 from utils import large
 from utils import logging_utils
 from utils import subprocess42
-from utils import tools
 import fake_swarming
 import task_runner
 
 CLIENT_DIR = os.path.normpath(
     os.path.join(test_env_bot_code.BOT_DIR, '..', '..', '..', 'client'))
 
 sys.path.insert(0, os.path.join(CLIENT_DIR, 'tests'))
 import isolateserver_mock
 
 
@@ skipping 36 matching lines @@
       file_path.rmtree(self.root_dir)
     except OSError:
       print >> sys.stderr, 'Failed to delete %s' % self.root_dir
     finally:
       super(TestTaskRunnerBase, self).tearDown()
 
   @classmethod
   def get_task_details(cls, *args, **kwargs):
     return task_runner.TaskDetails(get_manifest(*args, **kwargs))
 
-  def gen_requests(self, cost_usd=0., **kwargs):
+  def gen_requests(self, cost_usd=0., auth_headers=None, **kwargs):
     return [
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
-        self.get_check_first(cost_usd),
+        self.get_check_first(cost_usd, auth_headers=auth_headers),
         {'ok': True},
       ),
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
-        self.get_check_final(**kwargs),
+        self.get_check_final(auth_headers=auth_headers, **kwargs),
         {'ok': True},
       ),
     ]
 
   def requests(self, **kwargs):
     """Generates the expected HTTP requests for a task run."""
     self.expected_requests(self.gen_requests(**kwargs))
 
-  def get_check_first(self, cost_usd):
+  def get_check_first(self, cost_usd, auth_headers=None):
     def check_first(kwargs):
       self.assertLessEqual(cost_usd, kwargs['data'].pop('cost_usd'))
       self.assertEqual(
         {
           'data': {
             'id': 'localhost',
             'task_id': 23,
           },
+          'follow_redirects': False,
+          'headers': auth_headers or {},
         },
         kwargs)
     return check_first
 
 
 class TestTaskRunner(TestTaskRunnerBase):
   def setUp(self):
     super(TestTaskRunner, self).setUp()
     self.mock(time, 'time', lambda: 1000000000.)
 
-  def get_check_final(self, exit_code=0, output_re=r'^hi\n$', outputs_ref=None):
+  def get_check_final(
+      self, exit_code=0, output_re=r'^hi\n$', outputs_ref=None,
+      auth_headers=None):
     def check_final(kwargs):
       # Ignore these values.
       kwargs['data'].pop('bot_overhead', None)
       kwargs['data'].pop('duration', None)
 
       output = ''
       if 'output' in kwargs['data']:
         output = base64.b64decode(kwargs['data'].pop('output'))
       self.assertTrue(re.match(output_re, output))
 
       expected = {
         'data': {
           'cost_usd': 10.,
           'exit_code': exit_code,
           'hard_timeout': False,
           'id': 'localhost',
           'io_timeout': False,
           'output_chunk_start': 0,
           'task_id': 23,
         },
+        'follow_redirects': False,
+        'headers': auth_headers or {},
       }
       if outputs_ref:
         expected['data']['outputs_ref'] = outputs_ref
       self.assertEqual(expected, kwargs)
     return check_final
 
-  def _run_command(self, task_details):
+  def _run_command(self, task_details, auth_headers_file=None):
     start = time.time()
     self.mock(time, 'time', lambda: start + 10)
     server = 'https://localhost:1'
     return task_runner.run_command(
-        server, task_details, self.work_dir, 3600., start, 1)
+        server, task_details, self.work_dir, auth_headers_file, 3600., start, 1)
 
   def test_load_and_run_raw(self):
     server = 'https://localhost:1'
 
     def run_command(
-        swarming_server, task_details, work_dir, cost_usd_hour, start,
-        min_free_space):
+        swarming_server, task_details, work_dir, auth_headers_file,
+        cost_usd_hour, start, min_free_space):
       self.assertEqual(server, swarming_server)
+      self.assertTrue(isinstance(task_details, task_runner.TaskDetails))
       # Necessary for OSX.
       self.assertEqual(
           os.path.realpath(self.work_dir), os.path.realpath(work_dir))
-      self.assertTrue(isinstance(task_details, task_runner.TaskDetails))
+      self.assertEqual(None, auth_headers_file)
       self.assertEqual(3600., cost_usd_hour)
       self.assertEqual(time.time(), start)
       self.assertEqual(1, min_free_space)
       return {
         u'exit_code': 1,
         u'hard_timeout': False,
         u'io_timeout': False,
         u'must_signal_internal_failure': None,
         u'version': task_runner.OUT_VERSION,
       }
     self.mock(task_runner, 'run_command', run_command)
 
     manifest = os.path.join(self.root_dir, 'manifest')
     with open(manifest, 'wb') as f:
       data = {
         'bot_id': 'localhost',
         'command': ['a'],
         'env': {'d': 'e'},
         'extra_args': [],
         'grace_period': 30.,
         'hard_timeout': 10,
         'io_timeout': 11,
         'isolated': None,
         'task_id': 23,
       }
       json.dump(data, f)
 
     out_file = os.path.join(self.root_dir, 'work', 'task_runner_out.json')
-    task_runner.load_and_run(manifest, server, 3600., time.time(), out_file, 1)
+    task_runner.load_and_run(
+        manifest, server, None, 3600., time.time(), out_file, 1)
     expected = {
       u'exit_code': 1,
       u'hard_timeout': False,
       u'io_timeout': False,
       u'must_signal_internal_failure': None,
       u'version': task_runner.OUT_VERSION,
     }
     with open(out_file, 'rb') as f:
       self.assertEqual(expected, json.load(f))
 
   def test_load_and_run_isolated(self):
     self.expected_requests([])
     server = 'https://localhost:1'
 
     def run_command(
-        swarming_server, task_details, work_dir, cost_usd_hour, start,
-        min_free_space):
+        swarming_server, task_details, work_dir, auth_headers_file,
+        cost_usd_hour, start, min_free_space):
       self.assertEqual(server, swarming_server)
+      self.assertTrue(isinstance(task_details, task_runner.TaskDetails))
       # Necessary for OSX.
       self.assertEqual(
           os.path.realpath(self.work_dir), os.path.realpath(work_dir))
-      self.assertTrue(isinstance(task_details, task_runner.TaskDetails))
+      self.assertEqual(None, auth_headers_file)
       self.assertEqual(3600., cost_usd_hour)
       self.assertEqual(time.time(), start)
       self.assertEqual(1, min_free_space)
       return {
         u'exit_code': 0,
         u'hard_timeout': False,
         u'io_timeout': False,
         u'must_signal_internal_failure': None,
         u'version': task_runner.OUT_VERSION,
       }
@@ skipping 12 matching lines @@
         'isolated': {
           'input': '123',
           'server': 'http://localhost:1',
           'namespace': 'default-gzip',
         },
         'task_id': 23,
       }
       json.dump(data, f)
 
     out_file = os.path.join(self.root_dir, 'work', 'task_runner_out.json')
-    task_runner.load_and_run(manifest, server, 3600., time.time(), out_file, 1)
+    task_runner.load_and_run(
+        manifest, server, None, 3600., time.time(), out_file, 1)
     expected = {
       u'exit_code': 0,
       u'hard_timeout': False,
       u'io_timeout': False,
       u'must_signal_internal_failure': None,
       u'version': task_runner.OUT_VERSION,
     }
     with open(out_file, 'rb') as f:
       self.assertEqual(expected, json.load(f))
 
   def test_run_command_raw(self):
     # This runs the command for real.
     self.requests(cost_usd=1, exit_code=0)
     task_details = self.get_task_details('print(\'hi\')')
     expected = {
       u'exit_code': 0,
       u'hard_timeout': False,
       u'io_timeout': False,
       u'must_signal_internal_failure': None,
       u'version': task_runner.OUT_VERSION,
     }
     self.assertEqual(expected, self._run_command(task_details))
 
+  def test_run_command_raw_with_auth(self):
+    headers_file = os.path.join(self.root_dir, 'header.json')
+    with open(headers_file, 'wb') as f:
+      json.dump({'A': 'a'}, f)
+
+    # This runs the command for real.
+    self.requests(cost_usd=1, exit_code=0, auth_headers={'A': 'a'})
+    task_details = self.get_task_details('print(\'hi\')')
+    expected = {
+      u'exit_code': 0,
+      u'hard_timeout': False,
+      u'io_timeout': False,
+      u'must_signal_internal_failure': None,
+      u'version': task_runner.OUT_VERSION,
+    }
+    self.assertEqual(
+        expected, self._run_command(
+            task_details, auth_headers_file=headers_file))
+
   def test_run_command_isolated(self):
     # This runs the command for real.
     self.requests(
         cost_usd=1, exit_code=0,
         outputs_ref={
           u'isolated': u'123',
           u'isolatedserver': u'http://localhost:1',
           u'namespace': u'default-gzip',
         })
     task_details = self.get_task_details(isolated={
@@ skipping 123 matching lines @@
               'cost_usd': 10.,
               'duration': 0.,
               'exit_code': 0,
               'hard_timeout': False,
               'id': 'localhost',
               'io_timeout': False,
               'output': base64.b64encode('hi!\n'),
               'output_chunk_start': 100002*4,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': {},
           },
           kwargs)
 
     requests = [
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
         {
           'data': {
             'cost_usd': 10.,
             'id': 'localhost',
             'task_id': 23,
           },
+          'follow_redirects': False,
+          'headers': {},
         },
         {'ok': True},
       ),
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
         {
           'data': {
             'cost_usd': 10.,
             'id': 'localhost',
             'output': base64.b64encode('hi!\n' * 100002),
             'output_chunk_start': 0,
             'task_id': 23,
           },
+          'follow_redirects': False,
+          'headers': {},
         },
         {'ok': True},
       ),
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
         check_final,
         {'ok': True},
       ),
     ]
     self.expected_requests(requests)
@@ skipping 13 matching lines @@
       u'exit_code': 0,
       u'hard_timeout': False,
       u'io_timeout': False,
       u'must_signal_internal_failure': None,
       u'version': task_runner.OUT_VERSION,
     }
     self.assertEqual(expected, self._run_command(task_details))
 
   def test_main(self):
     def load_and_run(
-        manifest, swarming_server, cost_usd_hour, start, json_file,
-        min_free_space):
+        manifest, swarming_server, auth_headers_file, cost_usd_hour, start,
+        json_file, min_free_space):
       self.assertEqual('foo', manifest)
       self.assertEqual('http://localhost', swarming_server)
+      self.assertEqual(None, auth_headers_file)
       self.assertEqual(3600., cost_usd_hour)
       self.assertEqual(time.time(), start)
       self.assertEqual('task_summary.json', json_file)
       self.assertEqual(1, min_free_space)
 
     self.mock(task_runner, 'load_and_run', load_and_run)
     cmd = [
       '--swarming-server', 'http://localhost',
       '--in-file', 'foo',
       '--out-file', 'task_summary.json',
       '--cost-usd-hour', '3600',
       '--start', str(time.time()),
       '--min-free-space', '1',
     ]
     self.assertEqual(0, task_runner.main(cmd))
 
   def test_main_reboot(self):
     def load_and_run(
-        manifest, swarming_server, cost_usd_hour, start, json_file,
-        min_free_space):
+        manifest, swarming_server, auth_headers_file, cost_usd_hour, start,
+        json_file, min_free_space):
       self.assertEqual('foo', manifest)
       self.assertEqual('http://localhost', swarming_server)
+      self.assertEqual(None, auth_headers_file)
       self.assertEqual(3600., cost_usd_hour)
       self.assertEqual(time.time(), start)
       self.assertEqual('task_summary.json', json_file)
       self.assertEqual(1, min_free_space)
 
     self.mock(task_runner, 'load_and_run', load_and_run)
     cmd = [
       '--swarming-server', 'http://localhost',
       '--in-file', 'foo',
       '--out-file', 'task_summary.json',
@@ skipping 48 matching lines @@
     '  except IOError:\n'
     '    pass;\n'
     'print(\'bye\');\n'
     'time.sleep(100)') % (
         'signal.SIGBREAK' if sys.platform == 'win32' else 'signal.SIGTERM')
 
   SCRIPT_HANG = 'import time; print(\'hi\'); time.sleep(100)'
 
   def get_check_final(
       self, hard_timeout=False, io_timeout=False, exit_code=None,
-      output_re='^hi\n$'):
+      output_re='^hi\n$', auth_headers=None):
     def check_final(kwargs):
       kwargs['data'].pop('bot_overhead', None)
       if hard_timeout or io_timeout:
         self.assertLess(self.SHORT_TIME_OUT, kwargs['data'].pop('cost_usd'))
         self.assertLess(self.SHORT_TIME_OUT, kwargs['data'].pop('duration'))
       else:
         self.assertLess(0., kwargs['data'].pop('cost_usd'))
         self.assertLess(0., kwargs['data'].pop('duration'))
 
       output = ''
       if 'output' in kwargs['data']:
         output = base64.b64decode(kwargs['data'].pop('output'))
       self.assertTrue(re.match(output_re, output))
 
       self.assertEqual(
           {
             'data': {
               'exit_code': exit_code,
               'hard_timeout': hard_timeout,
               'id': 'localhost',
               'io_timeout': io_timeout,
               'output_chunk_start': 0,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': auth_headers or {},
           },
           kwargs)
     return check_final
 
   def _load_and_run(self, manifest):
     # Dot not mock time since this test class is testing timeouts.
     server = 'https://localhost:1'
     in_file = os.path.join(self.work_dir, 'task_runner_in.json')
     with open(in_file, 'wb') as f:
       json.dump(manifest, f)
     out_file = os.path.join(self.work_dir, 'task_runner_out.json')
-    task_runner.load_and_run(in_file, server, 3600., time.time(), out_file, 1)
+    task_runner.load_and_run(
+        in_file, server, None, 3600., time.time(), out_file, 1)
     with open(out_file, 'rb') as f:
       return json.load(f)
 
   def _run_command(self, task_details):
     # Dot not mock time since this test class is testing timeouts.
     server = 'https://localhost:1'
     return task_runner.run_command(
-        server, task_details, self.work_dir, 3600., time.time(), 1)
+        server, task_details, self.work_dir, None, 3600., time.time(), 1)
 
   def test_hard(self):
     # Actually 0xc000013a
     sig = -1073741510 if sys.platform == 'win32' else -signal.SIGTERM
     self.requests(hard_timeout=True, exit_code=sig)
     task_details = self.get_task_details(
         self.SCRIPT_HANG, hard_timeout=self.SHORT_TIME_OUT)
     expected = {
       u'exit_code': sig,
       u'hard_timeout': True,
@@ skipping 168 matching lines @@
                 },
                 u'upload': {
                   u'items_cold': [],
                   u'items_hot': [],
                 },
               },
               'output': 'hi\n',
               'output_chunk_start': 0,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': {},
           },
           kwargs)
     requests = [
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
         self.get_check_first(0.),
         {'ok': True},
       ),
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
@@ skipping 99 matching lines @@
                   u'items_hot': [],
                 },
                 u'upload': {
                   u'items_cold': [],
                   u'items_hot': [],
                 },
               },
               'output_chunk_start': 0,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': {},
           },
           kwargs)
     requests = [
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
         self.get_check_first(0.),
         {'ok': True},
       ),
       (
         'https://localhost:1/swarming/api/v1/bot/task_update/23',
@@ skipping 163 matching lines @@
   fix_encoding.fix_encoding()
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
   logging_utils.prepare_logging(None)
   logging_utils.set_console_level(
       logging.DEBUG if '-v' in sys.argv else logging.CRITICAL+1)
   # Fix litteral text expectation.
   os.environ['LANG'] = 'en_US.UTF-8'
   os.environ['LANGUAGE'] = 'en_US.UTF-8'
   unittest.main()