def main(args):
parser = optparse.OptionParser(description=sys.modules[__name__].__doc__)
parser.add_option('--in-file', help='Name of the request file')
parser.add_option('--out-file',
help='Name of the JSON file to write a task summary to')
parser.add_option('--swarming-server',
help='Swarming server to send data back')
parser.add_option('--cost-usd-hour',
type='float',
help='Cost of this VM in $/h')
parser.add_option('--start',
type='float',
help='Time this task was started')
options, args = parser.parse_args(args)
if not options.in_file or not options.out_file or args:
parser.error('task_runner is meant to be used by swarming_bot.')
on_error.report_on_exception_exit(options.swarming_server)
logging.info('starting')
remote = xsrf_client.XsrfRemote(options.swarming_server)
now = monotonic_time()
if options.start > now:
options.start = now
try:
load_and_run(options.in_file, remote, options.cost_usd_hour,
options.start, options.out_file)
return 0
finally:
logging.info('quitting')
def test_load_and_run_raw(self):
requests = [
(
'https://localhost:1/f',
{},
compress_to_zip({'file3': 'content3'}),
None,
),
]
self.expected_requests(requests)
server = xsrf_client.XsrfRemote('https://localhost:1/')
def run_command(swarming_server, task_details, work_dir, cost_usd_hour,
start):
self.assertEqual(server, swarming_server)
# Necessary for OSX.
self.assertEqual(os.path.realpath(self.work_dir), work_dir)
self.assertTrue(isinstance(task_details, task_runner.TaskDetails))
self.assertEqual(3600., cost_usd_hour)
self.assertEqual(time.time(), start)
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')
with open(manifest, 'wb') as f:
data = {
'bot_id': 'localhost',
'command': ['a'],
'data': [('https://localhost:1/f', 'foo.zip')],
'env': {
'd': 'e'
},
'extra_args': [],
'grace_period': 30.,
'hard_timeout': 10,
'inputs_ref': None,
'io_timeout': 11,
'task_id': 23,
}
json.dump(data, f)
out_file = os.path.join(self.root_dir, 'task_runner_out.json')
task_runner.load_and_run(manifest, server, 3600., time.time(),
out_file)
expected = {
u'exit_code': 1,
u'hard_timeout': False,
u'io_timeout': False,
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 setUp(self):
super(TestBotMain, self).setUp()
os.environ.pop('SWARMING_LOAD_TEST', None)
self.root_dir = tempfile.mkdtemp(prefix='bot_main')
self.old_cwd = os.getcwd()
os.chdir(self.root_dir)
# __main__ does it for us.
os.mkdir('logs')
self.server = xsrf_client.XsrfRemote('https://localhost:1/')
self.attributes = {
'dimensions': {
'foo': ['bar'],
'id': ['localhost'],
},
'state': {
'cost_usd_hour': 3600.,
},
'version': '123',
}
self.mock(zip_package, 'generate_version', lambda: '123')
self.bot = bot.Bot(self.server, self.attributes,
'https://localhost:1/', 'version1', self.root_dir,
self.fail)
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 _run_command(self, task_details):
start = time.time()
self.mock(time, 'time', lambda: start + 10)
server = xsrf_client.XsrfRemote('https://localhost:1/')
return task_runner.run_command(
server, task_details, '.', 3600., start,
os.path.join(self.work_dir, 'task_summary.json'))
def testXsrfRemoteSimple(self):
self.expected_requests([
(
'http://localhost/auth/api/v1/accounts/self/xsrf_token',
{
'data': {},
'headers': {
'X-XSRF-Token-Request': '1'
}
},
{
'xsrf_token': 'token'
},
),
(
'http://localhost/a',
{
'data': {
'foo': 'bar'
},
'headers': {
'X-XSRF-Token': 'token'
}
},
'foo',
None,
),
])
remote = xsrf_client.XsrfRemote('http://localhost/')
self.assertEqual('foo', remote.url_read('/a', data={'foo': 'bar'}))
def _load_and_run(self, manifest):
# Dot not mock time since this test class is testing timeouts.
server = xsrf_client.XsrfRemote('https://localhost:1/')
with open('manifest.json', 'w') as f:
json.dump(manifest, f)
return task_runner.load_and_run(
'manifest.json', server, 3600., time.time(),
os.path.join(self.work_dir, 'task_summary.json'))
def testXsrfRemoteRefresh(self):
self.expected_requests([
(
'http://localhost/auth/api/v1/accounts/self/xsrf_token',
{
'data': {},
'headers': {
'X-XSRF-Token-Request': '1'
}
},
{
'xsrf_token': 'token'
},
),
(
'http://localhost/a',
{
'data': {
'foo': 'bar'
},
'headers': {
'X-XSRF-Token': 'token'
}
},
# Fake that the token went bad by returning None. XsrfRemote will
# automatically try to refresh the token before retrying.
None,
None,
),
(
'http://localhost/auth/api/v1/accounts/self/xsrf_token',
{
'data': {},
'headers': {
'X-XSRF-Token-Request': '1'
}
},
{
'xsrf_token': 'token2'
},
),
(
'http://localhost/a',
{
'data': {
'foo': 'bar'
},
'headers': {
'X-XSRF-Token': 'token2'
}
},
'foo',
None,
),
])
remote = xsrf_client.XsrfRemote('http://localhost/')
remote.url_read('/a', data={'foo': 'bar'})
def get_remote():
"""Return a XsrfRemote instance to the preconfigured server."""
global _ERROR_HANDLER_WAS_REGISTERED
config = get_config()
server = config['server']
if not _ERROR_HANDLER_WAS_REGISTERED:
on_error.report_on_exception_exit(server)
_ERROR_HANDLER_WAS_REGISTERED = True
return xsrf_client.XsrfRemote(server, '/swarming/api/v1/bot/handshake')
def test_load_and_run_fail(self):
requests = [
(
'https://localhost:1/f',
{},
compress_to_zip({'file3': 'content3'}),
None,
),
]
self.expected_requests(requests)
server = xsrf_client.XsrfRemote('https://localhost:1/')
runs = []
def run_command(swarming_server, task_details, work_dir, cost_usd_hour,
start, json_file):
self.assertEqual(server, swarming_server)
# Necessary for OSX.
self.assertEqual(os.path.realpath(self.work_dir), work_dir)
self.assertTrue(isinstance(task_details, task_runner.TaskDetails))
self.assertEqual(3600., cost_usd_hour)
self.assertEqual(time.time(), start)
self.assertEqual('task_summary.json', json_file)
runs.append(0)
# Fails the first, pass the second.
return 1 if len(runs) == 1 else 0
self.mock(task_runner, 'run_command', run_command)
manifest = os.path.join(self.root_dir, 'manifest')
with open(manifest, 'wb') as f:
data = {
'bot_id': 'localhost',
'command': ['a'],
'data': [('https://localhost:1/f', 'foo.zip')],
'env': {
'd': 'e'
},
'grace_period': 30.,
'hard_timeout': 10,
'io_timeout': 11,
'task_id': 23,
}
json.dump(data, f)
self.assertEqual(
False,
task_runner.load_and_run(manifest, server, 3600., time.time(),
'task_summary.json'))
self.assertEqual([0], runs)
def testXsrfRemoteCustom(self):
# Use the new swarming bot API as an example of custom XSRF request handler.
self.expected_requests([
(
'http://localhost/swarming/api/v1/bot/handshake',
{
'data': {
'attributes': 'b'
},
'headers': {
'X-XSRF-Token-Request': '1'
},
},
{
'expiration_sec': 100,
'ignored': True,
'xsrf_token': 'token',
},
),
(
'http://localhost/a',
{
'data': {
'foo': 'bar'
},
'headers': {
'X-XSRF-Token': 'token'
}
},
'foo',
None,
),
])
remote = xsrf_client.XsrfRemote('http://localhost/',
'/swarming/api/v1/bot/handshake')
remote.xsrf_request_params = {'attributes': 'b'}
self.assertEqual('foo', remote.url_read('/a', data={'foo': 'bar'}))
def main(args):
parser = optparse.OptionParser(description=sys.modules[__name__].__doc__,
version=__version__)
parser.add_option('--file', help='Name of the request file')
parser.add_option('--swarming-server',
help='Swarming server to send data back')
parser.add_option('--cost-usd-hour',
type='float',
help='Cost of this VM in $/h')
parser.add_option('--start',
type='float',
help='Time this task was started')
parser.add_option('--json-file',
help='Name of the JSON file to write a task summary to')
options, args = parser.parse_args(args)
if not options.file:
parser.error('You must provide the request file name.')
if args:
parser.error('Unknown args: %s' % args)
on_error.report_on_exception_exit(options.swarming_server)
logging.info('starting')
remote = xsrf_client.XsrfRemote(options.swarming_server)
now = monotonic_time()
if options.start > now:
options.start = now
try:
if not load_and_run(options.file, remote, options.cost_usd_hour,
options.start, options.json_file):
return TASK_FAILED
return 0
finally:
logging.info('quitting')
def setUp(self):
super(TestBotMain, self).setUp()
os.environ.pop('SWARMING_LOAD_TEST', None)
self.root_dir = tempfile.mkdtemp(prefix='bot_main')
self.old_cwd = os.getcwd()
os.chdir(self.root_dir)
self.server = xsrf_client.XsrfRemote('https://localhost:1/')
self.attributes = {
'dimensions': {
'foo': ['bar'],
'id': ['localhost'],
},
'state': {
'cost_usd_hour': 3600.,
},
'version': '123',
}
self.mock(zip_package, 'generate_version', lambda: '123')
self.bot = bot.Bot(self.server, self.attributes, 'version1',
self.root_dir, self.fail)
self.mock(self.bot, 'post_error', self.fail)
self.mock(self.bot, 'restart', self.fail)
self.mock(subprocess, 'call', self.fail)
self.mock(time, 'time', lambda: 100.)
def _run_command(self, task_details):
# Dot not mock time since this test class is testing timeouts.
server = xsrf_client.XsrfRemote('https://localhost:1/')
return task_runner.run_command(server, task_details, self.work_dir,
3600., time.time())
def test_run_command_large(self):
# Method should have "self" as first argument - pylint: disable=E0213
class Popen(object):
"""Mocks the process so we can control how data is returned."""
def __init__(self2, cmd, cwd, env, stdout, stderr, stdin,
detached):
self.assertEqual(task_details.command, cmd)
self.assertEqual('./', cwd)
expected_env = os.environ.copy()
expected_env['foo'] = 'bar'
self.assertEqual(expected_env, env)
self.assertEqual(subprocess.PIPE, stdout)
self.assertEqual(subprocess.STDOUT, stderr)
self.assertEqual(subprocess.PIPE, stdin)
self.assertEqual(True, detached)
self2._out = [
'hi!\n',
'hi!\n',
'hi!\n' * 100000,
'hi!\n',
]
def yield_any(self2, maxsize, soft_timeout):
self.assertLess(0, maxsize)
self.assertLess(0, soft_timeout)
for i in self2._out:
yield 'stdout', i
@staticmethod
def wait():
return 0
@staticmethod
def kill():
self.fail()
self.mock(subprocess42, 'Popen', Popen)
def check_final(kwargs):
self.assertEqual(
{
'data': {
# That's because the cost includes the duration starting at start,
# not when the process was started.
'cost_usd': 10.,
'duration': 0.,
'exit_code': 0,
'hard_timeout': False,
'id': 'localhost',
'io_timeout': False,
'output': base64.b64encode('hi!\n'),
'output_chunk_start': 100002 * 4,
'task_id': 23,
},
'headers': {
'X-XSRF-Token': 'token'
},
},
kwargs)
requests = [
(
'https://localhost:1/auth/api/v1/accounts/self/xsrf_token',
{
'data': {},
'headers': {
'X-XSRF-Token-Request': '1'
}
},
{
'xsrf_token': 'token'
},
),
(
'https://localhost:1/swarming/api/v1/bot/task_update/23',
{
'data': {
'cost_usd': 10.,
'id': 'localhost',
'task_id': 23,
},
'headers': {
'X-XSRF-Token': 'token'
},
},
{},
),
(
'https://localhost:1/swarming/api/v1/bot/task_update/23',
{
'data': {
'cost_usd': 10.,
'id': 'localhost',
'output': base64.b64encode('hi!\n' * 100002),
'output_chunk_start': 0,
'task_id': 23,
},
'headers': {
'X-XSRF-Token': 'token'
},
},
{},
),
(
'https://localhost:1/swarming/api/v1/bot/task_update/23',
check_final,
{},
),
]
self.expected_requests(requests)
server = xsrf_client.XsrfRemote('https://localhost:1/')
task_details = task_runner.TaskDetails({
'bot_id':
'localhost',
'command': ['large', 'executable'],
'data': [],
'env': {
'foo': 'bar'
},
'grace_period':
30.,
'hard_timeout':
60,
'io_timeout':
60,
'task_id':
23,
})
start = time.time()
self.mock(time, 'time', lambda: start + 10)
r = task_runner.run_command(
server, task_details, './', 3600., start,
os.path.join(self.work_dir, 'task_summary.json'))
self.assertEqual(0, r)
def testXsrfRemoteGET(self):
self.expected_requests([('http://localhost/a', {}, 'foo', None)])
remote = xsrf_client.XsrfRemote('http://localhost/')
self.assertEqual('foo', remote.url_read('/a'))
def testXsrfRemoteRefresh(self):
self.expected_requests([
(
'http://localhost/auth/api/v1/accounts/self/xsrf_token',
{
'data': {},
'headers': {
'X-XSRF-Token-Request': '1'
}
},
{
'expiration_sec': 100,
'xsrf_token': 'token',
},
),
(
'http://localhost/a',
{
'data': {
'foo': 'bar'
},
'headers': {
'X-XSRF-Token': 'token'
}
},
'bar',
None,
),
(
'http://localhost/auth/api/v1/accounts/self/xsrf_token',
{
'data': {},
'headers': {
'X-XSRF-Token-Request': '1'
}
},
{
'expiration_sec': 100,
'xsrf_token': 'token2',
},
),
(
'http://localhost/a',
{
'data': {
'foo': 'bar'
},
'headers': {
'X-XSRF-Token': 'token2'
}
},
'foo',
None,
),
])
now = xsrf_client._utcnow()
remote = xsrf_client.XsrfRemote('http://localhost/')
remote.url_read('/a', data={'foo': 'bar'})
self.mock(xsrf_client, '_utcnow',
lambda: now + datetime.timedelta(seconds=91))
remote.url_read('/a', data={'foo': 'bar'})
