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