Send authorization headers when calling Swarming backend.

appengine/swarming/swarming_bot/bot_code/bot_main_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 json
 import logging
 import os
-import shutil
 import sys
 import tempfile
 import threading
 import time
 import unittest
 import zipfile
 
 import test_env_bot_code
 test_env_bot_code.setup_test_env()
 
 # Creates a server mock for functions in net.py.
 import net_utils
 
 import bot_main
+import remote_client
 from api import bot
 from api import os_utilities
 from depot_tools import fix_encoding
 from utils import file_path
 from utils import logging_utils
 from utils import net
 from utils import subprocess42
 from utils import zip_package
 
 
@@ skipping 18 matching lines @@
         'id': ['localhost'],
         'pool': ['default'],
       },
       'state': {
         'cost_usd_hour': 3600.,
         'sleep_streak': 0,
       },
       'version': '123',
     }
     self.mock(zip_package, 'generate_version', lambda: '123')
-    self.bot = bot.Bot(
-        self.attributes, 'https://localhost:1', 'version1', self.root_dir,
-        self.fail)
+    self.bot = self.make_bot()
     self.mock(self.bot, 'post_error', self.fail)
     self.mock(self.bot, 'restart', self.fail)
     self.mock(subprocess42, 'call', self.fail)
     self.mock(time, 'time', lambda: 100.)
     config_path = os.path.join(
         test_env_bot_code.BOT_DIR, 'config', 'config.json')
     with open(config_path, 'rb') as f:
       config = json.load(f)
     self.mock(bot_main, 'get_config', lambda: config)
     self.mock(
         bot_main, 'THIS_FILE',
         os.path.join(test_env_bot_code.BOT_DIR, 'swarming_bot.zip'))
 
   def tearDown(self):
     os.environ.pop('SWARMING_BOT_ID', None)
     os.chdir(self.old_cwd)
     file_path.rmtree(self.root_dir)
     super(TestBotMain, self).tearDown()
 
+  def make_bot(self, auth_headers_cb=None):
+    return bot.Bot(
+        remote_client.RemoteClient('https://localhost:1', auth_headers_cb),
+        self.attributes, 'https://localhost:1', 'version1',
+        self.root_dir, self.fail)
+
   def test_get_dimensions(self):
     dimensions = set(bot_main.get_dimensions(None))
     dimensions.discard('hidpi')
     dimensions.discard('zone')  # Only set on GCE bots.
     expected = {'cores', 'cpu', 'gpu', 'id', 'machine_type', 'os', 'pool'}
     if sys.platform == 'darwin':
       expected.add('xcode_version')
     self.assertEqual(expected, dimensions)
 
   def test_get_dimensions_load_test(self):
@@ skipping 62 matching lines @@
     self.expected_requests(
         [
           (
             'https://localhost:1/swarming/api/v1/bot/task_error/23',
             {
               'data': {
                 'id': expected_attribs['dimensions']['id'][0],
                 'message': 'error',
                 'task_id': 23,
               },
+              'follow_redirects': False,
+              'headers': {},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
             },
             {'resp': 1},
           ),
         ])
     botobj = bot_main.get_bot()
     self.assertEqual({'resp': 1}, bot_main.post_error_task(botobj, 'error', 23))
 
   def test_run_bot(self):
     # Test the run_bot() loop. Does not use self.bot.
     self.mock(time, 'time', lambda: 126.0)
@@ skipping 44 matching lines @@
     self.mock(subprocess42, 'Popen', Popen)
 
     self.expected_requests(
         [
           (
             'https://localhost:1/swarming/api/v1/bot/server_ping',
             {}, 'foo', None,
           ),
           (
             'https://localhost:1/swarming/api/v1/bot/handshake',
-            {'data': self.attributes},
+            {
+              'data': self.attributes,
+              'follow_redirects': False,
+              'headers': {},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
+            },
             {'bot_version': '123', 'server': self.url, 'server_version': 1},
           ),
         ])
 
     with self.assertRaises(Foo):
       bot_main.run_bot(None)
     self.assertEqual(
         self.attributes['dimensions']['id'][0], os.environ['SWARMING_BOT_ID'])
 
   def test_poll_server_sleep(self):
     slept = []
     bit = threading.Event()
     self.mock(bit, 'wait', slept.append)
     self.mock(bot_main, 'run_manifest', self.fail)
     self.mock(bot_main, 'update_bot', self.fail)
 
     self.expected_requests(
         [
           (
             'https://localhost:1/swarming/api/v1/bot/poll',
-            {'data': self.attributes},
+            {
+              'data': self.attributes,
+              'follow_redirects': False,
+              'headers': {},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
+            },
             {
               'cmd': 'sleep',
               'duration': 1.24,
+            },
+          ),
+        ])
+    self.assertFalse(bot_main.poll_server(self.bot, bit))
+    self.assertEqual([1.24], slept)
+
+  def test_poll_server_sleep_with_auth(self):
+    slept = []
+    bit = threading.Event()
+    self.mock(bit, 'wait', slept.append)
+    self.mock(bot_main, 'run_manifest', self.fail)
+    self.mock(bot_main, 'update_bot', self.fail)
+
+    self.bot = self.make_bot(lambda: ({'A': 'a'}, time.time() + 3600))
+
+    self.expected_requests(
+        [
+          (
+            'https://localhost:1/swarming/api/v1/bot/poll',
+            {
+              'data': self.attributes,
+              'follow_redirects': False,
+              'headers': {'A': 'a'},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
+            },
+            {
+              'cmd': 'sleep',
+              'duration': 1.24,
             },
           ),
         ])
     self.assertFalse(bot_main.poll_server(self.bot, bit))
     self.assertEqual([1.24], slept)
 
   def test_poll_server_run(self):
     manifest = []
     clean = []
     bit = threading.Event()
     self.mock(bit, 'wait', self.fail)
     self.mock(bot_main, 'run_manifest', lambda *args: manifest.append(args))
     self.mock(bot_main, 'clean_isolated_cache',
               lambda *args: clean.append(args))
     self.mock(bot_main, 'update_bot', self.fail)
 
     self.expected_requests(
         [
           (
             'https://localhost:1/swarming/api/v1/bot/poll',
-            {'data': self.bot._attributes},
+            {
+              'data': self.bot._attributes,
+              'follow_redirects': False,
+              'headers': {},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
+            },
             {
               'cmd': 'run',
               'manifest': {'foo': 'bar'},
             },
           ),
         ])
     self.assertTrue(bot_main.poll_server(self.bot, bit))
     expected = [(self.bot, {'foo': 'bar'}, time.time())]
     self.assertEqual(expected, manifest)
     expected = [(self.bot,)]
     self.assertEqual(expected, clean)
 
   def test_poll_server_update(self):
     update = []
     bit = threading.Event()
     self.mock(bit, 'wait', self.fail)
     self.mock(bot_main, 'run_manifest', self.fail)
     self.mock(bot_main, 'update_bot', lambda *args: update.append(args))
 
     self.expected_requests(
         [
           (
             'https://localhost:1/swarming/api/v1/bot/poll',
-            {'data': self.attributes},
+            {
+              'data': self.attributes,
+              'follow_redirects': False,
+              'headers': {},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
+            },
             {
               'cmd': 'update',
               'version': '123',
             },
           ),
         ])
     self.assertTrue(bot_main.poll_server(self.bot, bit))
     self.assertEqual([(self.bot, '123')], update)
 
   def test_poll_server_restart(self):
     restart = []
     bit = threading.Event()
     self.mock(bit, 'wait', self.fail)
     self.mock(bot_main, 'run_manifest', self.fail)
     self.mock(bot_main, 'update_bot', self.fail)
     self.mock(self.bot, 'restart', lambda *args: restart.append(args))
 
     self.expected_requests(
         [
           (
             'https://localhost:1/swarming/api/v1/bot/poll',
-            {'data': self.attributes},
+            {
+              'data': self.attributes,
+              'follow_redirects': False,
+              'headers': {},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
+            },
             {
               'cmd': 'restart',
               'message': 'Please die now',
             },
           ),
         ])
     self.assertTrue(bot_main.poll_server(self.bot, bit))
     self.assertEqual([('Please die now',)], restart)
 
   def test_poll_server_restart_load_test(self):
     os.environ['SWARMING_LOAD_TEST'] = '1'
     bit = threading.Event()
     self.mock(bit, 'wait', self.fail)
     self.mock(bot_main, 'run_manifest', self.fail)
     self.mock(bot_main, 'update_bot', self.fail)
     self.mock(self.bot, 'restart', self.fail)
 
     self.expected_requests(
         [
           (
             'https://localhost:1/swarming/api/v1/bot/poll',
             {
               'data': self.attributes,
+              'follow_redirects': False,
+              'headers': {},
+              'timeout': remote_client.NET_CONNECTION_TIMEOUT_SEC,
             },
             {
               'cmd': 'restart',
               'message': 'Please die now',
             },
           ),
         ])
     self.assertTrue(bot_main.poll_server(self.bot, bit))
 
-  def _mock_popen(self, returncode=0, exit_code=0, url='https://localhost:1'):
+  def _mock_popen(
+      self, returncode=0, exit_code=0, url='https://localhost:1',
+      auth_params_json=None):
     result = {
       'exit_code': exit_code,
       'must_signal_internal_failure': None,
       'version': 3,
     }
     # Method should have "self" as first argument - pylint: disable=E0213
     class Popen(object):
       def __init__(
           self2, cmd, detached, cwd, env, stdout, stderr, stdin, close_fds):
         self2.returncode = None
@@ skipping 12 matching lines @@
             1024 * 1024)),
         ]
         self.assertEqual(expected, cmd)
         self.assertEqual(True, detached)
         self.assertEqual(self.bot.base_dir, cwd)
         self.assertEqual('24', env['SWARMING_TASK_ID'])
         self.assertTrue(stdout)
         self.assertEqual(subprocess42.STDOUT, stderr)
         self.assertEqual(subprocess42.PIPE, stdin)
         self.assertEqual(sys.platform != 'win32', close_fds)
+        if auth_params_json is not None:
+          auth_params_file = os.path.join(
+              self.root_dir, 'work', 'bot_auth_params.json')
+          self.assertEqual(auth_params_file, env.get('SWARMING_AUTH_PARAMS'))
+          with open(auth_params_file, 'rb') as f:
+            actual_auth_params = json.load(f)
+          self.assertEqual(auth_params_json, actual_auth_params)
 
       def wait(self2, timeout=None): # pylint: disable=unused-argument
         self2.returncode = returncode
         with open(self2._out_file, 'wb') as f:
           json.dump(result, f)
         return 0
 
     self.mock(subprocess42, 'Popen', Popen)
     return result
 
@@ skipping 14 matching lines @@
       'command': ['echo', 'hi'],
       'dimensions': {'os': 'Amiga', 'pool': 'default'},
       'grace_period': 30,
       'hard_timeout': 60,
       'host': 'https://localhost:3',
       'task_id': '24',
     }
     self.assertEqual(self.root_dir, self.bot.base_dir)
     bot_main.run_manifest(self.bot, manifest, time.time())
 
+  def test_run_manifest_with_auth(self):
+    self.bot = self.make_bot(
+        auth_headers_cb=lambda: ({'A': 'a'}, time.time() + 3600))
+
+    self.mock(bot_main, 'post_error_task', lambda *args: self.fail(args))
+    def call_hook(botobj, name, *args):
+      if name == 'on_after_task':
+        failure, internal_failure, dimensions, summary = args
+        self.assertEqual(self.attributes['dimensions'], botobj.dimensions)
+        self.assertEqual(False, failure)
+        self.assertEqual(False, internal_failure)
+        self.assertEqual({'os': 'Amiga', 'pool': 'default'}, dimensions)
+        self.assertEqual(result, summary)
+    self.mock(bot_main, 'call_hook', call_hook)
+    result = self._mock_popen(
+        url='https://localhost:3',
+        auth_params_json={'swarming_http_headers': {'A': 'a'}})
+
+    manifest = {
+      'command': ['echo', 'hi'],
+      'dimensions': {'os': 'Amiga', 'pool': 'default'},
+      'grace_period': 30,
+      'hard_timeout': 60,
+      'host': 'https://localhost:3',
+      'task_id': '24',
+    }
+    self.assertEqual(self.root_dir, self.bot.base_dir)
+    bot_main.run_manifest(self.bot, manifest, time.time())
+
   def test_run_manifest_task_failure(self):
     self.mock(bot_main, 'post_error_task', lambda *args: self.fail(args))
     def call_hook(_botobj, name, *args):
       if name == 'on_after_task':
         failure, internal_failure, dimensions, summary = args
         self.assertEqual(True, failure)
         self.assertEqual(False, internal_failure)
         self.assertEqual({'pool': 'default'}, dimensions)
         self.assertEqual(result, summary)
     self.mock(bot_main, 'call_hook', call_hook)
@@ skipping 67 matching lines @@
     # Under the test however forever-blocking calls finish, and post_error is
     # called.
     self.mock(self.bot, 'post_error', lambda *_: None)
     # Mock the file to download in the temporary directory.
     self.mock(
         bot_main, 'THIS_FILE',
         os.path.join(self.root_dir, 'swarming_bot.1.zip'))
     new_zip = os.path.join(self.root_dir, 'swarming_bot.2.zip')
     # This is necessary otherwise zipfile will crash.
     self.mock(time, 'time', lambda: 1400000000)
-    def url_retrieve(f, url):
+    def url_retrieve(f, url, headers=None, timeout=None):
       self.assertEqual(
           'https://localhost:1/swarming/api/v1/bot/bot_code/123', url)
       self.assertEqual(new_zip, f)
+      self.assertEqual({}, headers)
+      self.assertEqual(remote_client.NET_CONNECTION_TIMEOUT_SEC, timeout)
       # Create a valid zip that runs properly.
       with zipfile.ZipFile(f, 'w') as z:
         z.writestr('__main__.py', 'print("hi")')
       return True
     self.mock(net, 'url_retrieve', url_retrieve)
 
     calls = []
     def exec_python(args):
       calls.append(args)
       return 23
@@ skipping 26 matching lines @@
     self.assertEqual(0, bot_main.main([]))
 
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   if '-v' in sys.argv:
     TestBotMain.maxDiff = None
   logging.basicConfig(
       level=logging.DEBUG if '-v' in sys.argv else logging.CRITICAL)
   unittest.main()