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 setUp(self):
api_test.ApiTest.setUp(self)
request_key = ndb.Key(
datastore_entities.Request, REQUEST_ID,
namespace=common.NAMESPACE)
self.request = datastore_test_util.CreateRequest(
request_id=request_key.id(),
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line',
cluster='cluster',
run_target='run_target')
])
self.request.put()
self.command = self._AddCommand(
REQUEST_ID,
'1',
command_line='cmd',
cluster='example_cluster',
run_target='hammerhead',
run_count=1,
ants_invocation_id='i123',
ants_work_unit_id='w123')
self.plugin_patcher = mock.patch(
'__main__.env_config.CONFIG.plugin')
self.plugin_patcher.start()
self.host = datastore_test_util.CreateHost('cluster', 'hostname', 'alab')
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 testSyncRequest_processError(self, mock_queue, mock_process):
request_sync_monitor.Monitor(REQUEST_ID)
sync_key = request_sync_monitor.GetRequestSyncStatusKey(REQUEST_ID)
sync_status = sync_key.get()
sync_status.has_new_command_events = True
sync_status.put()
mock_process.side_effect = RuntimeError
datastore_test_util.CreateRequest(request_id=REQUEST_ID,
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line2',
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.QUEUED)
with self.assertRaises(RuntimeError):
request_sync_monitor.SyncRequest(REQUEST_ID)
mock_process.assert_called_once_with(REQUEST_ID)
mock_queue.assert_called_once_with(REQUEST_ID) # Call from Monitor()
sync_status = sync_key.get()
self.assertTrue(sync_status.has_new_command_events)
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 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 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 setUp(self):
super(CommanderTest, self).setUp()
datastore_test_util.CreateRequest(user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line",
cluster="cluster",
run_target="run_target")
],
request_id=REQUEST_ID)
def testProcessRequests_missingRunTarget(self, cancel_request):
request_id = "1001"
datastore_test_util.CreateRequest(
user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line0 --shard-count 1",
cluster="cluster",
run_target=None)
],
request_id=request_id)
commander._ProcessRequest(request_id)
cancel_request.assert_called_once_with(
request_id, common.CancelReason.INVALID_REQUEST)
def _CreateTestRequest(self, state=common.RequestState.UNKNOWN):
"""Creates a Request for testing purposes."""
request = datastore_test_util.CreateRequest(
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line --request-id %d' % self._request_id,
cluster='cluster',
run_target='run_target')
],
request_id=str(self._request_id))
request.state = state
request.put()
self._request_id += 1
return request
def testNotifyNoDirtyRequest(self):
request = datastore_test_util.CreateRequest(
request_id="1",
user="******",
command_infos=[
datastore_entities.CommandInfo(command_line="command_line",
cluster="cluster",
run_target="run_target")
],
state=common.RequestState.UNKNOWN,
notify_state_change=False)
request.put()
notification_handler.NotifyRequestState(request_id="1")
tasks = self.mock_task_scheduler.GetTasks()
self.assertEqual(0, len(tasks))
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 testSyncRequest(self, mock_update, mock_queue, mock_process):
mock_update.return_value = True
datastore_test_util.CreateRequest(request_id=REQUEST_ID,
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line2',
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.QUEUED)
request_sync_monitor.SyncRequest(REQUEST_ID)
mock_process.assert_called_once_with(REQUEST_ID)
mock_queue.assert_called_once_with(
REQUEST_ID,
countdown_secs=request_sync_monitor.SHORT_SYNC_COUNTDOWN_SECS)
def testCheckPendingCommands_canceledRequest(self, 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(10)
]
request = datastore_test_util.CreateRequest(
request_id=request_id,
user="******",
command_infos=command_infos,
max_concurrent_tasks=5,
plugin_data={
"FOO": "foo",
"BAR": "'bar",
})
command_manager.CreateCommands(request_id=request_id,
command_infos=command_infos,
priority=request.priority,
shard_indexes=[0] * len(command_infos))
request.state = common.RequestState.CANCELED
request.put()
commands = command_manager.GetCommands(request_id)
for i, command in enumerate(commands):
if i < 2:
command.state = common.CommandState.COMPLETED
elif i < 5:
command.state = common.CommandState.QUEUED
else:
command.state = common.CommandState.UNKNOWN
command.put()
request_summary = request_manager.RequestSummary()
request_summary.completed_count = 2
request_summary.queued_count = 3
request_summary.pending_count = 5
commander._CheckPendingCommands(request, request_summary)
schedule_tasks.assert_not_called()
monitor.assert_not_called()
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 testNotifyNoDirtyRequest_force(self, mock_add_task):
request = datastore_test_util.CreateRequest(
request_id="1",
user="******",
command_infos=[
datastore_entities.CommandInfo(command_line="command_line",
cluster="cluster",
run_target="run_target")
],
state=common.RequestState.UNKNOWN,
notify_state_change=False)
request.put()
notification_handler.NotifyRequestState(request_id="1", force=True)
mock_add_task.assert_called_once_with(
queue_name=common.OBJECT_EVENT_QUEUE,
payload=mock.ANY,
transactional=True)
payload = zlib.decompress(mock_add_task.call_args[1]["payload"])
task = json.loads(payload)
self.assertEqual("1", task["request_id"])
def testSyncRequest_finalRequest(self, mock_update, mock_queue,
mock_process):
mock_update.return_value = True
request = datastore_test_util.CreateRequest(
request_id=REQUEST_ID,
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line2',
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.COMPLETED)
request.put()
request_sync_monitor.SyncRequest(REQUEST_ID)
sync_key = request_sync_monitor.GetRequestSyncStatusKey(REQUEST_ID)
self.assertIsNone(sync_key.get())
mock_process.assert_called_once_with(REQUEST_ID)
mock_queue.assert_not_called()
def testNotifyDirtyRequest_Error(self, add):
request = datastore_test_util.CreateRequest(
request_id="1",
user="******",
command_infos=[
datastore_entities.CommandInfo(command_line="command_line",
cluster="cluster",
run_target="run_target")
],
state=common.RequestState.UNKNOWN,
notify_state_change=True)
request.put()
# Make sure that notify fails.
add.side_effect = errors.Error()
try:
notification_handler.NotifyRequestState(request_id="1")
self.fail("apiclient.errors.Error should have been thrown")
except errors.Error:
# expected failure
pass
# Since the pub sub failed, we should have the dirty bit still set.
r = request_manager.GetRequest("1")
self.assertTrue(r.notify_state_change)
def setUp(self):
api_test.ApiTest.setUp(self)
request_key = ndb.Key(datastore_entities.Request,
'1001',
namespace=common.NAMESPACE)
request = datastore_test_util.CreateRequest(
request_id=request_key.id(),
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line',
cluster='cluster',
run_target='run_target')
])
request.put()
command = datastore_entities.Command(parent=request_key,
id='1',
command_line='command_line',
cluster='cluster',
run_target='run_target',
state=common.CommandState.RUNNING)
command.put()
command_attempt_0 = datastore_entities.CommandAttempt(
parent=command.key,
id='attempt_id-0',
task_id='task-id-0',
attempt_id='attempt_id-0',
hostname='hostname_0',
device_serial='device_serial_0',
create_time=TIME,
state=common.CommandState.COMPLETED)
command_attempt_0.put()
command_attempt_1 = datastore_entities.CommandAttempt(
parent=command.key,
id='attempt_id-1',
task_id='task-id-1',
attempt_id='attempt_id-1',
hostname='hostname_0',
device_serial='device_serial_1',
create_time=TIME + TIME_DELTA,
state=common.CommandState.RUNNING)
command_attempt_1.put()
command_attempt_2 = datastore_entities.CommandAttempt(
parent=command.key,
id='attempt_id-2',
task_id='task-id-2',
attempt_id='attempt_id-2',
hostname='hostname_1',
device_serial='device_serial_2',
create_time=TIME + TIME_DELTA * 2,
state=common.CommandState.RUNNING)
command_attempt_2.put()
command_attempt_3 = datastore_entities.CommandAttempt(
parent=command.key,
id='attempt_id-3',
task_id='task-id-3',
attempt_id='attempt_id-3',
hostname='hostname_1',
device_serial='device_serial_2',
create_time=TIME + TIME_DELTA * 3,
device_serials=['d2', 'd3'],
state=common.CommandState.RUNNING)
command_attempt_3.put()
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, {}),
])
def testProcessCommandEvent_pendingCommands(self, attempt_metric, monitor):
# Test ProcessCommandEvent for a non-final state with deletion
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(10)
]
request = datastore_test_util.CreateRequest(
request_id=request_id,
user="******",
command_infos=command_infos,
max_concurrent_tasks=5,
plugin_data={
"FOO": "foo",
"BAR": "'bar",
})
commands = command_manager.CreateCommands(request_id=request_id,
command_infos=command_infos,
priority=request.priority,
shard_indexes=[0] *
len(command_infos))
command_manager.ScheduleTasks(commands[:5])
_, request_id, _, command_id = commands[0].key.flat()
pending_commands = command_manager.GetCommands(
request_id, common.CommandState.UNKNOWN)
self.assertEqual(5, len(pending_commands))
queued_commands = command_manager.GetCommands(
request_id, common.CommandState.QUEUED)
self.assertEqual(5, len(queued_commands))
tasks = command_manager.GetActiveTasks(commands[0])
self.assertEqual(1, len(tasks))
command_task_store.LeaseTask(tasks[0].task_id)
command_event_test_util.CreateCommandAttempt(
commands[0],
"attempt0",
common.CommandState.UNKNOWN,
task=tasks[0])
event = command_event_test_util.CreateTestCommandEvent(
request_id,
command_id,
"attempt0",
common.InvocationEventType.INVOCATION_COMPLETED,
task=tasks[0],
time=TIMESTAMP)
commander.ProcessCommandEvent(event)
tasks = command_manager.GetActiveTasks(commands[0])
self.assertEqual(0, len(tasks))
command = commands[0].key.get(use_cache=False)
self.assertEqual(common.CommandState.COMPLETED, command.state)
attempt_metric.assert_called_once_with(cluster_id=command.cluster,
run_target=command.run_target,
hostname="hostname",
state="COMPLETED")
next_command = pending_commands[0]
monitor.assert_called_once_with([next_command])
next_command = pending_commands[0].key.get(use_cache=False)
self.assertEqual(common.CommandState.QUEUED, next_command.state)
pending_commands = command_manager.GetCommands(
request_id, common.CommandState.UNKNOWN)
self.assertEqual(4, len(pending_commands))
queued_commands = command_manager.GetCommands(
request_id, common.CommandState.QUEUED)
self.assertEqual(5, len(queued_commands))
completed_commands = command_manager.GetCommands(
request_id, common.CommandState.COMPLETED)
self.assertEqual(1, len(completed_commands))
