def testPost(self, schedule_tasks):
request_id = "request_id"
request = datastore_test_util.CreateRequest(
user="******",
command_infos=[
datastore_entities.CommandInfo(command_line="command_line0",
cluster="cluster",
run_target="bullhead")
],
request_id=request_id,
plugin_data={
"ants_invocation_id": "i123",
"ants_work_unit_id": "w123"
})
request_manager.AddToQueue(request)
tasks = self.mock_task_scheduler.GetTasks()
self.assertEqual(len(tasks), 1)
self.testapp.post(commander.REQUEST_HANDLER_PATH, tasks[0].payload)
commands = request_manager.GetCommands(request_id)
self.assertEqual(1, len(commands))
command = commands[0]
self.assertEqual("command_line0", command.command_line)
self.assertEqual("bullhead", command.run_target)
self.assertEqual(1, command.run_count)
self.assertEqual("cluster", command.cluster)
self.assertIsNone(command.priority)
self.assertIsNone(command.queue_timeout_seconds)
def testProcessRequest_escapeInCommandLine(self, plugin, schedule_tasks):
request_id1 = "1001"
datastore_test_util.CreateRequest(
user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line=("command_line0"
' --arg \'option=\'"\'"\'value\'"\'"\'\''),
cluster="cluster",
run_target="run_target")
],
request_id=request_id1)
commander._ProcessRequest(request_id1)
self.assertEqual(1, schedule_tasks.call_count)
commands = request_manager.GetCommands(request_id1)
self.assertEqual(1, len(commands))
command = commands[0]
self.assertEqual("command_line0 --arg option='value'",
command.command_line)
self.assertEqual("run_target", command.run_target)
self.assertEqual(1, command.run_count)
self.assertEqual("cluster", command.cluster)
plugin.assert_has_calls([])
def testProcessRequests_shardedRequests_oneShard(self, plugin,
schedule_tasks):
"""Tests processing of sharded requests with a single shard."""
request_id = "1001"
datastore_test_util.CreateRequest(user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line0",
cluster="cluster",
run_target="bullhead",
shard_count=1)
],
request_id=request_id)
commander._ProcessRequest(request_id)
commands = request_manager.GetCommands(request_id)
self.assertEqual(1, len(commands))
command = commands[0]
command_line = command_util.CommandLine(command.command_line)
command_line.RemoveOptions(["--shard-index"])
# If only one shard is specified, the --shard-count option is removed.
self.assertEqual("command_line0", command_line.ToTFString())
self.assertEqual("bullhead", command.run_target)
self.assertEqual(1, command.run_count)
self.assertEqual("cluster", command.cluster)
plugin.assert_has_calls([])
def testProcessRequests_RequestlocalSharding(self, plugin, schedule_tasks):
"""Tests processing of sharded requests with local sharding."""
request_id = "1001"
datastore_test_util.CreateRequest(
user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line0 --shard-count 3",
cluster="cluster",
run_target="bullhead")
],
request_id=request_id)
commander._ProcessRequest(request_id)
commands = request_manager.GetCommands(request_id)
self.assertEqual(1, len(commands))
for command in commands:
command_line = command_util.CommandLine(command.command_line)
command_line.RemoveOptions(["--shard-index"])
self.assertEqual("command_line0 --shard-count 3",
command_line.ToTFString())
self.assertEqual("bullhead", command.run_target)
self.assertEqual(1, command.run_count)
self.assertEqual("cluster", command.cluster)
plugin.assert_has_calls([])
def testProcessRequests_shardedRequests(self, plugin, schedule_tasks):
"""Tests processing of sharded requests."""
request_id = "1001"
datastore_test_util.CreateRequest(request_id=request_id,
user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line0",
cluster="cluster",
run_target="bullhead",
shard_count=3)
])
commander._ProcessRequest(request_id)
commands = request_manager.GetCommands(request_id)
self.assertEqual(3, len(commands))
shards = []
for command in commands:
command_line = command_util.CommandLine(command.command_line)
shards.append(command_line.GetOption("--shard-index"))
command_line.RemoveOptions(["--shard-index"])
self.assertEqual("command_line0 --shard-count 3",
command_line.ToTFString())
self.assertEqual("bullhead", command.run_target)
self.assertEqual(1, command.run_count)
self.assertEqual("cluster", command.cluster)
self.assertCountEqual(["0", "1", "2"], shards)
plugin.assert_has_calls([])
def testProcessRequest_withMultipleCommands(self, plugin, schedule_tasks,
monitor):
request_id = "1001"
command_infos = [
datastore_entities.CommandInfo(command_line="command_line %04d" %
i,
cluster="cluster %04d" % i,
run_target="run_target %04d" % i,
run_count=1,
shard_count=1) for i in range(500)
]
request = datastore_test_util.CreateRequest(
request_id=request_id,
user="******",
command_infos=command_infos,
max_concurrent_tasks=100,
plugin_data={
"FOO": "foo",
"BAR": "'bar",
})
commander._ProcessRequest(request_id)
commands = request_manager.GetCommands(request_id)
commands.sort(key=lambda x: x.command_line)
self.assertEqual(len(command_infos), len(commands))
for command_info, command in zip(command_infos, commands):
self.assertEqual(command_info.command_line, command.command_line)
self.assertEqual(command_info.cluster, command.cluster)
self.assertEqual(command_info.run_target, command.run_target)
self.assertEqual(command_info.run_count, command.run_count)
self.assertIsNone(command.shard_count)
plugin.OnCreateCommands.assert_has_calls([
mock.call([
plugin_base.CommandInfo(command_id=int(command.key.id()),
command_line=command.command_line,
run_count=command.run_count,
shard_count=1,
shard_index=0) for command in commands
], request.plugin_data, {}),
])
schedule_tasks.assert_called_once_with(
commands[:request.max_concurrent_tasks])
monitor.assert_called_once_with(
commands[:request.max_concurrent_tasks])
def testProcessRequest_invalidRequest(self, plugin, cancel_request,
schedule_tasks):
request_id1 = "1001"
datastore_test_util.CreateRequest(request_id=request_id1,
user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line0",
cluster="cluster",
run_target=None)
])
commander._ProcessRequest(request_id1)
self.assertFalse(schedule_tasks.called)
commands = request_manager.GetCommands(request_id1)
self.assertEqual(0, len(commands))
cancel_request.assert_called_once_with(
request_id1, common.CancelReason.INVALID_REQUEST)
plugin.assert_has_calls([])
def testProcessRequest(self, plugin, schedule_tasks, monitor):
request_id1 = "1001"
request_id2 = "1002"
datastore_test_util.CreateRequest(
user="******",
command_infos=[
datastore_entities.CommandInfo(command_line=(
"command_line0 --run-target run_target --cluster cluster"),
cluster="cluster",
run_target="run_target")
],
request_id=request_id1,
plugin_data={
"ants_invocation_id": "i123",
"ants_work_unit_id": "w123"
})
datastore_test_util.CreateRequest(
user="******",
command_infos=[
datastore_entities.CommandInfo(command_line=(
"command_line0 --run-target run_target --cluster cluster"),
cluster="cluster",
run_target="run_target")
],
request_id=request_id2,
priority=100,
queue_timeout_seconds=86400)
commander._ProcessRequest(request_id1)
commander._ProcessRequest(request_id2)
self.assertEqual(2, schedule_tasks.call_count)
self.assertEqual(2, monitor.call_count)
commands_0 = request_manager.GetCommands(request_id1)
self.assertEqual(1, len(commands_0))
command = commands_0[0]
self.assertEqual("command_line0", command.command_line)
self.assertEqual("run_target", command.run_target)
self.assertEqual(1, command.run_count)
self.assertEqual("cluster", command.cluster)
commands_1 = request_manager.GetCommands(request_id2)
self.assertEqual(1, len(commands_1))
command = commands_1[0]
self.assertEqual("command_line0", command.command_line)
self.assertEqual("run_target", command.run_target)
self.assertEqual(1, command.run_count)
self.assertEqual("cluster", command.cluster)
self.assertEqual(100, command.priority)
self.assertEqual(86400, command.queue_timeout_seconds)
monitor.assert_has_calls(
[mock.call(commands_0),
mock.call(commands_1)])
plugin.assert_has_calls([
mock.call.OnCreateCommands([
plugin_base.CommandInfo(command_id=5629499534213120,
command_line="command_line0",
run_count=1,
shard_count=1,
shard_index=0)
], {
"ants_invocation_id": "i123",
"ants_work_unit_id": "w123"
}, {}),
mock.call.OnCreateCommands([
plugin_base.CommandInfo(command_id=5066549580791808,
command_line="command_line0",
run_count=1,
shard_count=1,
shard_index=0)
], None, {}),
])