def testListScheduledFlows(self, db: abstract_db.Database): context = _CreateContext(db) client_id1 = db_test_utils.InitializeClient(db) client_id2 = db_test_utils.InitializeClient(db) handler = flow_plugin.ApiScheduleFlowHandler() sf1 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id1, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), context=context) sf2 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id1, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), context=context) handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id2, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), context=context) handler = flow_plugin.ApiListScheduledFlowsHandler() args = flow_plugin.ApiListScheduledFlowsArgs(client_id=client_id1, creator=context.username) results = handler.Handle(args, context=context) self.assertEqual(results.scheduled_flows, [sf1, sf2])
def testUnscheduleFlowRemovesScheduledFlow(self, db: abstract_db.Database): token = _CreateToken(db) client_id = db_test_utils.InitializeClient(db) handler = flow_plugin.ApiScheduleFlowHandler() sf1 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), token=token) sf2 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), token=token) handler = flow_plugin.ApiUnscheduleFlowHandler() args = flow_plugin.ApiUnscheduleFlowArgs( client_id=client_id, scheduled_flow_id=sf1.scheduled_flow_id) handler.Handle(args, token=token) handler = flow_plugin.ApiListScheduledFlowsHandler() args = flow_plugin.ApiListScheduledFlowsArgs(client_id=client_id, creator=token.username) results = handler.Handle(args, token=token) self.assertEqual(results.scheduled_flows, [sf2])
def _ListScheduledFlows(client_id: str, creator: str): handler = api_flow.ApiListScheduledFlowsHandler() return handler.Handle(api_flow.ApiListScheduledFlowsArgs( client_id=client_id, creator=creator), context=api_call_context.ApiCallContext( username=creator)).scheduled_flows