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))