def testScheduleFlowCreatesMultipleScheduledFlows(self):
client_id0 = self.SetupClient(0)
client_id1 = self.SetupClient(1)
username0 = self.SetupUser("u0")
username1 = self.SetupUser("u1")
sf0 = self.ScheduleFlow(client_id=client_id0, creator=username0)
sf1 = self.ScheduleFlow(client_id=client_id0, creator=username0)
sf2 = self.ScheduleFlow(client_id=client_id1, creator=username0)
sf3 = self.ScheduleFlow(client_id=client_id0, creator=username1)
self.assertEqual([sf0, sf1],
flow.ListScheduledFlows(client_id0, username0))
self.assertEqual([sf2], flow.ListScheduledFlows(client_id1, username0))
self.assertEqual([sf3], flow.ListScheduledFlows(client_id0, username1))
def testStartScheduledFlowsContinuesNextOnFailure(self):
client_id = self.SetupClient(0)
username = self.SetupUser("u0")
self.ScheduleFlow(
client_id=client_id,
creator=username,
flow_name=file.CollectSingleFile.__name__,
flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"),
runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))
self.ScheduleFlow(
client_id=client_id,
creator=username,
flow_name=file.CollectSingleFile.__name__,
flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"),
runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))
with mock.patch.object(
rdf_file_finder.CollectSingleFileArgs,
"Validate",
side_effect=[ValueError("foobazzle"), mock.DEFAULT]):
flow.StartScheduledFlows(client_id, username)
self.assertLen(data_store.REL_DB.ReadAllFlowObjects(client_id), 1)
self.assertLen(flow.ListScheduledFlows(client_id, username), 1)
def Handle(self, args, context=None):
results = flow.ListScheduledFlows(
client_id=args.client_id, creator=args.creator)
results = sorted(results, key=lambda sf: sf.create_time)
results = [ApiScheduledFlow.FromScheduledFlow(sf) for sf in results]
return ApiListScheduledFlowsResult(scheduled_flows=results)
def testStartScheduledFlowsDeletesScheduledFlows(self):
client_id = self.SetupClient(0)
username = self.SetupUser("u0")
self.ScheduleFlow(client_id=client_id, creator=username)
self.ScheduleFlow(client_id=client_id, creator=username)
flow.StartScheduledFlows(client_id, username)
self.assertEmpty(flow.ListScheduledFlows(client_id, username))
def testUnscheduleFlowCorrectlyRemovesScheduledFlow(self):
client_id = self.SetupClient(0)
username = self.SetupUser("u0")
sf1 = self.ScheduleFlow(client_id=client_id, creator=username)
sf2 = self.ScheduleFlow(client_id=client_id, creator=username)
flow.UnscheduleFlow(client_id, username, sf1.scheduled_flow_id)
self.assertEqual([sf2], flow.ListScheduledFlows(client_id, username))
flow.StartScheduledFlows(client_id, username)
self.assertLen(data_store.REL_DB.ReadAllFlowObjects(client_id), 1)
def testStartScheduledFlowsHandlesErrorInFlowArgsValidation(self):
client_id = self.SetupClient(0)
username = self.SetupUser("u0")
self.ScheduleFlow(
client_id=client_id,
creator=username,
flow_name=file.CollectSingleFile.__name__,
flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"),
runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))
with mock.patch.object(rdf_file_finder.CollectSingleFileArgs,
"Validate",
side_effect=ValueError("foobazzle")):
flow.StartScheduledFlows(client_id, username)
self.assertEmpty(data_store.REL_DB.ReadAllFlowObjects(client_id))
scheduled_flows = flow.ListScheduledFlows(client_id, username)
self.assertLen(scheduled_flows, 1)
self.assertIn("foobazzle", scheduled_flows[0].error)
