def setUp(self):
super(CommandEventHandlerTest, self).setUp()
self.testapp = webtest.TestApp(command_event_handler.APP)
self.plugin_patcher = mock.patch(
"__main__.env_config.CONFIG.plugin")
self.plugin_patcher.start()
self.request = request_manager.CreateRequest(
request_id="1001",
user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line",
cluster="cluster",
run_target="run_target"),
])
self.command = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line="long command line",
cluster="foobar",
run_target="foo",
run_count=1,
shard_count=1),
],
shard_indexes=list(range(1)),
request_plugin_data={
"ants_invocation_id": "i123",
"ants_work_unit_id": "w123"
})[0]
self.now_patcher = mock.patch.object(common, "Now")
self.mock_now = self.now_patcher.start()
self.mock_now.return_value = TIMESTAMP
def setUp(self):
super(CommandMonitorTest, self).setUp()
self.testapp = webtest.TestApp(command_monitor.APP)
self.plugin_patcher = mock.patch(
'__main__.env_config.CONFIG.plugin')
self.plugin_patcher.start()
self.request = request_manager.CreateRequest(
request_id='1001',
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line',
cluster='cluster',
run_target='run_target')
])
self.request_2 = request_manager.CreateRequest(
request_id='1002',
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line',
cluster='cluster',
run_target='run_target')
])
# Clear Datastore cache
ndb.get_context().clear_cache()
def setUp(self):
super(CommandAttemptMonitorTest, self).setUp()
self.testapp = webtest.TestApp(command_attempt_monitor.APP)
self.plugin_patcher = mock.patch(
'__main__.env_config.CONFIG.plugin')
self.plugin_patcher.start()
self.request = request_manager.CreateRequest(
request_id='1001',
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line',
cluster='cluster',
run_target='run_target'),
])
self.command = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='long command line',
cluster='foobar',
run_target='foo',
run_count=1,
shard_count=1),
],
shard_indexes=list(range(1)),
request_plugin_data={
'ants_invocation_id': 'i123',
'ants_work_unit_id': 'w123'
})[0]
# Clear Datastore cache
ndb.get_context().clear_cache()
def testMonitor(self, sync):
commands = command_manager.CreateCommands(
request_id=self.request_2.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='long command line',
shard_count=2,
run_target='foo',
run_count=1,
cluster='foobar'),
datastore_entities.CommandInfo(
command_line='longer_command_line',
shard_count=2,
run_target='foo',
run_count=1,
cluster='foobar'),
],
shard_indexes=list(range(2)))
num_monitored = command_monitor.Monitor(commands)
self.assertEqual(2, num_monitored)
tasks = self.mock_task_scheduler.GetTasks()
self.assertEqual(2, len(tasks))
response_0 = self.testapp.post(
'/_ah/queue/%s' % command_monitor.COMMAND_SYNC_QUEUE,
tasks[0].payload)
self.assertEqual('200 OK', response_0.status)
response_1 = self.testapp.post(
'/_ah/queue/%s' % command_monitor.COMMAND_SYNC_QUEUE,
tasks[1].payload)
self.assertEqual('200 OK', response_1.status)
sync.assert_has_calls([
mock.call(self.request_2.key.id(), commands[0].key.id()),
mock.call(self.request_2.key.id(), commands[1].key.id())
])
def testEnqueueCommandEvents_multipleEvents(self):
self.request = request_manager.CreateRequest(
request_id="9999",
user="******",
command_infos=[
datastore_entities.CommandInfo(
command_line="command_line",
cluster="cluster",
run_target="run_target",
shard_count=2)
])
command_1, command_2 = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line="long command line 0",
cluster="foobar",
run_target="foo",
run_count=1,
shard_count=2),
datastore_entities.CommandInfo(
command_line="long command line 1",
cluster="foobar",
run_target="foo",
run_count=1,
shard_count=2)
],
shard_indexes=list(range(2)))
_, request_id, _, command_1_id = command_1.key.flat()
_, _, _, command_2_id = command_2.key.flat()
command_event_test_util.CreateCommandAttempt(
command_1, "aid", common.CommandState.QUEUED)
command_event_test_util.CreateCommandAttempt(
command_2, "aid", common.CommandState.QUEUED)
event = command_event_test_util.CreateTestCommandEventJson(
request_id, command_1_id, "aid", "InvocationStarted")
event2 = command_event_test_util.CreateTestCommandEventJson(
request_id, command_2_id, "aid", "InvocationStarted")
event3 = command_event_test_util.CreateTestCommandEventJson(
request_id, command_1_id, "aid", "InvocationCompleted")
event4 = command_event_test_util.CreateTestCommandEventJson(
request_id, command_2_id, "aid", "InvocationCompleted")
command_event_handler.EnqueueCommandEvents([event, event2, event3, event4])
tasks = self.mock_task_scheduler.GetTasks()
self.assertEqual(len(tasks), 4)
for task in tasks:
self.testapp.post(
command_event_handler.COMMAND_EVENT_HANDLER_PATH, task.payload)
command_attempts = command_manager.GetCommandAttempts(
request_id, command_1_id)
self.assertEqual(len(command_attempts), 1)
self.assertEqual(common.CommandState.COMPLETED, command_attempts[0].state)
command_attempts = command_manager.GetCommandAttempts(
request_id, command_2_id)
self.assertEqual(len(command_attempts), 1)
self.assertEqual(common.CommandState.COMPLETED, command_attempts[0].state)
def setUp(self):
api_test.ApiTest.setUp(self)
self.request1 = request_manager.CreateRequest(
request_id='1',
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line1',
cluster='cluster',
run_target='run_target')
])
self.request1.state = common.RequestState.RUNNING
self.request1.put()
self.request2 = request_manager.CreateRequest(
request_id='2',
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line2',
cluster='cluster',
run_target='run_target')
])
self.request2.state = common.RequestState.COMPLETED
self.request2.put()
self.requests = [self.request1, self.request2]
self.command1 = datastore_entities.Command(
parent=self.request1.key,
id='1',
request_id='1',
command_line='command_line',
cluster='cluster',
run_target='run_target',
run_count=1,
state=common.CommandState.RUNNING,
start_time=START_TIME,
end_time=None,
create_time=START_TIME,
update_time=START_TIME)
self.command1.put()
self.command2 = datastore_entities.Command(
parent=self.request1.key,
id='2',
request_id='1',
command_line='command_line',
cluster='cluster',
run_target='run_target',
run_count=1,
state=common.CommandState.RUNNING,
start_time=START_TIME,
end_time=None,
create_time=START_TIME,
update_time=START_TIME)
self.command2.put()
self.commands = [self.command1, self.command2]
def testSyncCommand_withCustomQueueTimeout(self, mock_ensure, sync):
datastore_entities.Command.update_time._auto_now = False
now = datetime.datetime.utcnow()
command_1, command_2 = command_manager.CreateCommands(
request_id=self.request_2.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='long command line',
shard_count=2,
run_target='foo',
run_count=1,
cluster='foobar'),
datastore_entities.CommandInfo(
command_line='longer_command_line',
shard_count=2,
run_target='foo',
run_count=1,
cluster='foobar'),
],
shard_indexes=list(range(2)),
queue_timeout_seconds=command_monitor.MAX_COMMAND_INACTIVE_TIME_MIN *
2 * 60)
# Change update times. command_1 should ensure leasable, command_2 should
# ensure leasable and cancel afterwards
command_1.state = common.CommandState.QUEUED
command_1.update_time = (
now - datetime.timedelta(
minutes=command_monitor.MAX_COMMAND_INACTIVE_TIME_MIN))
command_1.put()
command_2.state = common.CommandState.QUEUED
command_2.update_time = (
now - datetime.timedelta(
minutes=command_monitor.MAX_COMMAND_INACTIVE_TIME_MIN) * 3)
command_2.put()
command_monitor.SyncCommand(command_1.request_id, command_1.key.id())
command_monitor.SyncCommand(command_2.request_id, command_2.key.id())
mock_ensure.assert_has_calls([
mock.call(
hamcrest.match_equality(
hamcrest.has_property('key', command_1.key))),
mock.call(
hamcrest.match_equality(
hamcrest.has_property('key', command_2.key)))
])
self.assertEqual(common.CommandState.QUEUED, command_1.key.get().state)
self.assertEqual(common.CommandState.CANCELED, command_2.key.get().state)
self.assertEqual(common.RequestState.CANCELED,
self.request_2.key.get().state)
sync.assert_called_once_with(command_1)
def testAddToSyncQueue_RunningCommand(self, mock_add, mock_now):
# Create a command that has been running for 3 hours.
datastore_entities.Command.update_time._auto_now = False
now = datetime.datetime.utcnow()
mock_now.return_value = now
command = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='command line',
run_target='run_target',
run_count=1,
shard_count=1,
cluster='cluster')
],
shard_indexes=list(range(1)))[0]
_, request_id, _, command_id = command.key.flat()
command.state = common.CommandState.RUNNING
command.update_time = now - datetime.timedelta(hours=3)
command.put()
command_monitor.AddToSyncQueue(command)
# Command monitor should schedule it to be synced in
# MAX_COMMAND_EVENT_DELAY_MINs.
payload = json.dumps({
command_manager.COMMAND_ID_KEY: command_id,
command_manager.REQUEST_ID_KEY: request_id,
})
mock_add.assert_called_once_with(
queue_name=command_monitor.COMMAND_SYNC_QUEUE,
payload=payload,
eta=now + datetime.timedelta(
minutes=command_monitor.MAX_COMMAND_EVENT_DELAY_MIN))
def testAddToSyncQueue_CustomCancelDeadline(self, mock_add):
# Create a command with a custom 10 hour command timeout that needs to be
# cancelled in 1 minute.
datastore_entities.Command.update_time._auto_now = False
now = datetime.datetime.utcnow()
custom_timeout = 10 * 3600
command = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='command line',
run_target='run_target',
run_count=1,
shard_count=1,
cluster='cluster')
],
shard_indexes=list(range(1)),
queue_timeout_seconds=custom_timeout)[0]
_, request_id, _, command_id = command.key.flat()
command.state = common.CommandState.QUEUED
command.update_time = now - datetime.timedelta(seconds=custom_timeout - 60)
command.put()
command_monitor.AddToSyncQueue(command)
# Command monitor should schedule it to be synced in 1 minute.
payload = json.dumps({
command_manager.COMMAND_ID_KEY: command_id,
command_manager.REQUEST_ID_KEY: request_id,
})
mock_add.assert_called_once_with(
queue_name=command_monitor.COMMAND_SYNC_QUEUE,
payload=payload,
eta=now + datetime.timedelta(minutes=1))
def testSyncCommand(self, mock_ensure, sync):
datastore_entities.Command.update_time._auto_now = False
now = datetime.datetime.utcnow()
command = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='long command line',
cluster='foobar',
run_target='foo',
run_count=1,
shard_count=1)
],
shard_indexes=list(range(1)))[0]
command.state = common.CommandState.QUEUED
command.update_time = (
now - datetime.timedelta(
minutes=command_monitor.MAX_COMMAND_INACTIVE_TIME_MIN) * 2)
command.put()
command_monitor.SyncCommand(command.request_id, command.key.id())
mock_ensure.assert_called_once_with(
hamcrest.match_equality(hamcrest.has_property('key', command.key)))
self.assertEqual(common.CommandState.CANCELED, command.key.get().state)
self.assertEqual(common.RequestState.CANCELED, self.request.key.get().state)
sync.assert_not_called()
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 _CreateAttempt(self, attempt_id, task_id, state):
# Helper to create an attempt
command = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='long command line',
run_target='foo',
run_count=1,
shard_count=1,
cluster='foobar')
],
shard_indexes=list(range(1)),
request_plugin_data={
'ants_invocation_id': 'i123',
'ants_work_unit_id': 'w123'
})[0]
_, request_id, _, command_id = command.key.flat()
attempt_key = ndb.Key(
datastore_entities.Request, request_id,
datastore_entities.Command, command_id,
datastore_entities.CommandAttempt, attempt_id,
namespace=common.NAMESPACE)
attempt = datastore_entities.CommandAttempt(
key=attempt_key,
attempt_id=attempt_id,
state=state,
command_id=command_id,
task_id=task_id)
attempt.put()
return attempt
def testListRequests(self):
for request_id in range(1, 11):
request_id = str(request_id)
request_manager.CreateRequest(
user='******',
request_id=request_id,
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line1',
cluster='cluster',
run_target='run_target')
])
api_request = {
'user': '******',
'state': 0,
'offset': 1,
'count': 2,
}
api_response = self.testapp.post_json('/_ah/api/RequestApi.ListRequest',
api_request)
request_collection = (protojson
.decode_message(api_messages.RequestMessageCollection,
api_response.body))
self.assertEqual(2, len(request_collection.requests))
self.assertEqual('9', request_collection.requests[0].id)
self.assertEqual('8', request_collection.requests[1].id)
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 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 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 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 testBackfillCommands(self, mock_add):
command_1, command_2, command_3 = command_manager.CreateCommands(
request_id=self.request.key.id(),
command_infos=[
datastore_entities.CommandInfo(
command_line='long command line',
shard_count=3,
run_target='foo',
run_count=1,
cluster='foobar'),
datastore_entities.CommandInfo(
command_line='longer_command_line',
shard_count=3,
run_target='foo',
run_count=1,
cluster='foobar'),
datastore_entities.CommandInfo(
command_line='short_cmd',
shard_count=3,
run_target='foo',
run_count=1,
cluster='foobar'),
],
shard_indexes=list(range(3)),
request_plugin_data={
'ants_invocation_id': 'i123',
'ants_work_unit_id': 'w123'
})
command_1.state = common.CommandState.QUEUED
command_1.put()
command_2.state = common.CommandState.QUEUED
command_2.put()
command_3.state = common.CommandState.RUNNING
command_3.put()
response = self.testapp.post_json(
'/_ah/api/CoordinatorApi.BackfillCommands', {})
self.assertEqual('200 OK', response.status)
mock_add.assert_has_calls(
[
mock.call(
hamcrest.match_equality(
hamcrest.has_property('key', command_1.key))),
mock.call(
hamcrest.match_equality(
hamcrest.has_property('key', command_2.key))),
], any_order=True)
self.assertEqual(2, mock_add.call_count)
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 testBackfillRequestSyncs(self, mock_monitor):
queued_request = request_manager.CreateRequest(
request_id='queued_id', user='******', command_infos=[
datastore_entities.CommandInfo(command_line='command_line2')
])
queued_request.state = common.RequestState.QUEUED
queued_request.put()
response = self.testapp.post_json(
'/_ah/api/CoordinatorApi.BackfillRequestSyncs', {})
self.assertEqual('200 OK', response.status)
mock_monitor.assert_has_calls(
[mock.call(self.request.key.id()),
mock.call(queued_request.key.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 testRequestFromEntity(self):
"""Tests converting a Request entity to a message."""
request = datastore_entities.Request(
id='id',
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line',
shard_count=2,
run_count=3,
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.UNKNOWN)
message = datastore_entities.ToMessage(request)
self.AssertEqualRequest(request, message)
def setUp(self):
super(CoordinatorApiTest, self).setUp()
self.plugin_patcher = mock.patch(
'__main__.env_config.CONFIG.plugin')
self.plugin_patcher.start()
self.request = request_manager.CreateRequest(
request_id='1001',
user='******',
command_infos=[
datastore_entities.CommandInfo(
command_line='command_line',
cluster='cluster',
run_target='run_target')
])
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 testRequestFromEntity_canceledRequest(self):
"""Tests converting a cancelled Request entity to a message."""
request = datastore_entities.Request(
id='id',
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line',
shard_count=2,
run_count=3,
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.CANCELED,
cancel_reason=common.CancelReason.QUEUE_TIMEOUT)
message = datastore_entities.ToMessage(request)
self.AssertEqualRequest(request, message)
def testMessageFromEntity_cancelRequest_invalidRequest(self):
"""Tests converting a Request object to a message."""
request_key = ndb.Key(datastore_entities.Request, '123')
request = datastore_entities.Request(
key=request_key,
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line',
shard_count=2,
run_count=3,
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.UNKNOWN,
cancel_reason=common.CancelReason.INVALID_REQUEST)
message = datastore_entities.ToMessage(request)
self.AssertEqualRequest(request, message)
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])
