Minor refactoring in task_runner.py in preparation for adding more code.

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 depot_tools import fix_encoding
 from utils import file_path
 from utils import large
 from utils import logging_utils
 from utils import subprocess42
+import bot_auth
 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
 
 
 def get_manifest(script=None, isolated=None, **kwargs):
   isolated_input = isolated and isolated.get('input')
   out = {
     'bot_id': 'localhost',
     'command':
         [sys.executable, '-u', '-c', script] if not isolated_input else None,
     'env': {},
     'extra_args': [],
     'grace_period': 30.,
     'hard_timeout': 10.,
     'io_timeout': 10.,
     'isolated': isolated,
     'task_id': 23,
   }
   out.update(kwargs)
   return out
 
 
+class FakeAuthSystem(object):
+  def __init__(self):
+    self._running = False
+
+  def start(self, auth_params_file):
+    assert auth_params_file == '/path/to/auth-params-file'
+    assert not self._running
+    self._running = True
+
+  def stop(self):
+    self._running = False
+
+  @property
+  def bot_headers(self):
+    assert self._running
+    return {'Fake': 'Header'}
+
+
 class TestTaskRunnerBase(net_utils.TestCase):
   def setUp(self):
     super(TestTaskRunnerBase, self).setUp()
     self.root_dir = tempfile.mkdtemp(prefix='task_runner')
     logging.info('Temp: %s', self.root_dir)
     self.work_dir = os.path.join(self.root_dir, 'w')
     os.chdir(self.root_dir)
     os.mkdir(self.work_dir)
     # Create the logs directory so run_isolated.py can put its log there.
     self._logs = os.path.join(self.root_dir, 'logs')
@@ skipping 88 matching lines @@
           '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, auth_params_file=None):
+  def _run_command(self, task_details, headers_cb=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,
-        '/path/to/file', auth_params_file)
+        '/path/to/file', headers_cb or (lambda: {}))
 
   def test_load_and_run_raw(self):
     server = 'https://localhost:1'
 
+    self.mock(bot_auth, 'AuthSystem', FakeAuthSystem)
+
     def run_command(
         swarming_server, task_details, work_dir,
-        cost_usd_hour, start, min_free_space, bot_file, auth_params_file):
+        cost_usd_hour, start, min_free_space, bot_file, headers_cb):
       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.assertEqual(3600., cost_usd_hour)
       self.assertEqual(time.time(), start)
       self.assertEqual(1, min_free_space)
       self.assertEqual('/path/to/bot-file', bot_file)
-      self.assertEqual('/path/to/auth-params-file', auth_params_file)
+      self.assertEqual({'Fake': 'Header'}, headers_cb())
       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')
@@ skipping 22 matching lines @@
       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'
 
+    self.mock(bot_auth, 'AuthSystem', FakeAuthSystem)
+
     def run_command(
         swarming_server, task_details, work_dir,
-        cost_usd_hour, start, min_free_space, bot_file, auth_params_file):
+        cost_usd_hour, start, min_free_space, bot_file, headers_cb):
       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.assertEqual(3600., cost_usd_hour)
       self.assertEqual(time.time(), start)
       self.assertEqual(1, min_free_space)
       self.assertEqual('/path/to/bot-file', bot_file)
-      self.assertEqual('/path/to/auth-params-file', auth_params_file)
+      self.assertEqual({'Fake': 'Header'}, headers_cb())
       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,
       }
     self.mock(task_runner, 'run_command', run_command)
 
     manifest = os.path.join(self.root_dir, 'manifest')
@@ skipping 36 matching lines @@
     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):
-    auth_params_file = os.path.join(self.root_dir, 'auth_params.json')
-    with open(auth_params_file, 'wb') as f:
-      json.dump({'swarming_http_headers': {'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_params_file=auth_params_file))
+        self._run_command(task_details, headers_cb=lambda: {'A': 'a'}))
 
   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',
         })
@@ skipping 335 matching lines @@
     task_runner.load_and_run(
         in_file, server, 3600., time.time(), out_file, 1, None, None)
     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,
-        '/path/to/file', None)
+        '/path/to/file', lambda: {})
 
   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 474 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()