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 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 testScheduleFlow(self, db: abstract_db.Database):
token = _CreateToken(db)
client_id = db_test_utils.InitializeClient(db)
handler = flow_plugin.ApiScheduleFlowHandler()
args = 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)))
sf = handler.Handle(args, token=token)
self.assertEqual(sf.client_id, client_id)
self.assertEqual(sf.creator, token.username)
self.assertNotEmpty(sf.scheduled_flow_id)
self.assertEqual(sf.flow_name, file.CollectSingleFile.__name__)
self.assertEqual(sf.flow_args.path, "/foo")
self.assertEqual(sf.runner_args.cpu_limit, 60)
def ScheduleFlow(
self,
args: api_flow.ApiCreateFlowArgs,
context: Optional[api_call_context.ApiCallContext] = None,
) -> api_flow.ApiScheduleFlowHandler:
return api_flow.ApiScheduleFlowHandler()
def ScheduleFlow(
self,
args: api_flow.ApiCreateFlowArgs,
token: Optional[access_control.ACLToken] = None
) -> api_flow.ApiScheduleFlowHandler:
return api_flow.ApiScheduleFlowHandler()
