Chromium Code Reviews
chromiumcodereview-hr@appspot.gserviceaccount.com (chromiumcodereview-hr) | Please choose your nickname with Settings | Help | Chromium Project | Gerrit Changes | Sign out
(2021)

Unified Diff: appengine/swarming/swarming_bot/bot_code/task_runner_test.py

Issue 2024313003: Send authorization headers when calling Swarming backend. (Closed) Base URL: https://chromium.googlesource.com/external/github.com/luci/luci-py@master
Patch Set: rebase Created 4 years, 6 months ago
Use n/p to move between diff chunks; N/P to move between comments. Draft comments are only viewable by you.
Jump to:
View side-by-side diff with in-line comments
Download patch
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 = [

Powered by Google App Engine
This is Rietveld 408576698