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)