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 = [ |