Send authorization headers when calling Swarming backend.

appengine/swarming/swarming_bot/bot_code/task_runner_test.py

Index: appengine/swarming/swarming_bot/bot_code/task_runner_test.py
diff --git a/appengine/swarming/swarming_bot/bot_code/task_runner_test.py b/appengine/swarming/swarming_bot/bot_code/task_runner_test.py
index ad7c828cf0657ea6caef99982082bbc9778582c4..7f39535535db079018341863bb41222874656467 100755
--- a/appengine/swarming/swarming_bot/bot_code/task_runner_test.py
+++ b/appengine/swarming/swarming_bot/bot_code/task_runner_test.py
@@ -21,13 +21,11 @@ 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
 
@@ -84,16 +82,16 @@ class TestTaskRunnerBase(net_utils.TestCase):
   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},
       ),
     ]
@@ -102,7 +100,7 @@ class TestTaskRunnerBase(net_utils.TestCase):
     """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(
@@ -111,6 +109,8 @@ class TestTaskRunnerBase(net_utils.TestCase):
             'id': 'localhost',
             'task_id': 23,
           },
+          'follow_redirects': False,
+          'headers': auth_headers or {},
         },
         kwargs)
     return check_first
@@ -121,7 +121,9 @@ class TestTaskRunner(TestTaskRunnerBase):
     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)
@@ -142,30 +144,33 @@ class TestTaskRunner(TestTaskRunnerBase):
           '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)
@@ -194,7 +199,8 @@ class TestTaskRunner(TestTaskRunnerBase):
       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,
@@ -210,13 +216,14 @@ class TestTaskRunner(TestTaskRunnerBase):
     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)
@@ -249,7 +256,8 @@ class TestTaskRunner(TestTaskRunnerBase):
       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,
@@ -273,6 +281,25 @@ class TestTaskRunner(TestTaskRunnerBase):
     }
     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(
@@ -416,6 +443,8 @@ class TestTaskRunner(TestTaskRunnerBase):
               'output_chunk_start': 100002*4,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': {},
           },
           kwargs)
 
@@ -428,6 +457,8 @@ class TestTaskRunner(TestTaskRunnerBase):
             'id': 'localhost',
             'task_id': 23,
           },
+          'follow_redirects': False,
+          'headers': {},
         },
         {'ok': True},
       ),
@@ -441,6 +472,8 @@ class TestTaskRunner(TestTaskRunnerBase):
             'output_chunk_start': 0,
             'task_id': 23,
           },
+          'follow_redirects': False,
+          'headers': {},
         },
         {'ok': True},
       ),
@@ -474,10 +507,11 @@ class TestTaskRunner(TestTaskRunnerBase):
 
   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)
@@ -496,10 +530,11 @@ class TestTaskRunner(TestTaskRunnerBase):
 
   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)
@@ -568,7 +603,7 @@ class TestTaskRunnerNoTimeMock(TestTaskRunnerBase):
 
   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:
@@ -593,6 +628,8 @@ class TestTaskRunnerNoTimeMock(TestTaskRunnerBase):
               'output_chunk_start': 0,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': auth_headers or {},
           },
           kwargs)
     return check_final
@@ -604,7 +641,8 @@ class TestTaskRunnerNoTimeMock(TestTaskRunnerBase):
     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)
 
@@ -612,7 +650,7 @@ class TestTaskRunnerNoTimeMock(TestTaskRunnerBase):
     # 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
@@ -801,6 +839,8 @@ class TestTaskRunnerNoTimeMock(TestTaskRunnerBase):
               'output_chunk_start': 0,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': {},
           },
           kwargs)
     requests = [
@@ -920,6 +960,8 @@ class TestTaskRunnerNoTimeMock(TestTaskRunnerBase):
               'output_chunk_start': 0,
               'task_id': 23,
             },
+            'follow_redirects': False,
+            'headers': {},
           },
           kwargs)
     requests = [