def Run(self):
# Create a running mock refresh operation.
self.running_flow_urn = self.CreateRecursiveListFlow(
self.client_id, self.token)
# Create a mock refresh operation and complete it.
self.finished_flow_urn = self.CreateRecursiveListFlow(
self.client_id, self.token)
with aff4.FACTORY.Open(
self.finished_flow_urn,
aff4_type=flow.GRRFlow,
mode="rw",
token=self.token) as flow_obj:
flow_obj.GetRunner().Error("Fake error")
# Create an arbitrary flow to check on 404s.
self.non_refresh_flow_urn = flow.GRRFlow.StartFlow(
client_id=self.client_id,
flow_name=discovery.Interrogate.__name__,
token=self.token)
# Unkonwn flow ids should also cause 404s.
self.unknown_flow_id = "F:12345678"
# Check both operations.
self.Check(
"GetVfsRefreshOperationState",
args=vfs_plugin.ApiGetVfsRefreshOperationStateArgs(
client_id=self.client_id.Basename(),
operation_id=str(self.running_flow_urn)),
replace={
self.running_flow_urn.Basename(): "W:ABCDEF"
})
self.Check(
"GetVfsRefreshOperationState",
args=vfs_plugin.ApiGetVfsRefreshOperationStateArgs(
client_id=self.client_id.Basename(),
operation_id=str(self.finished_flow_urn)),
replace={
self.finished_flow_urn.Basename(): "W:ABCDEF"
})
self.Check(
"GetVfsRefreshOperationState",
args=vfs_plugin.ApiGetVfsRefreshOperationStateArgs(
client_id=self.client_id.Basename(),
operation_id=str(self.non_refresh_flow_urn)),
replace={
self.non_refresh_flow_urn.Basename(): "W:ABCDEF"
})
self.Check(
"GetVfsRefreshOperationState",
args=vfs_plugin.ApiGetVfsRefreshOperationStateArgs(
client_id=self.client_id.Basename(),
operation_id=str(self.unknown_flow_id)),
replace={
self.unknown_flow_id: "W:ABCDEF"
})
def testHandlerThrowsExceptionOnUnknownFlowId(self):
# Create args with an operation id not referencing any flow.
args = vfs_plugin.ApiGetVfsRefreshOperationStateArgs(
operation_id="F:12345678")
# Our mock flow can't be read, so an error should be raised.
with self.assertRaises(vfs_plugin.VfsFileContentUpdateNotFoundError):
self.handler.Handle(args, token=self.token)
def testHandlerThrowsExceptionOnArbitraryFlowId(self):
# Create a mock flow.
self.flow_urn = flow.GRRFlow.StartFlow(
client_id=self.client_id, flow_name="Interrogate", token=self.token)
args = vfs_plugin.ApiGetVfsRefreshOperationStateArgs(
operation_id=str(self.flow_urn))
# Our mock flow is not a RecursiveListFlow, so an error should be raised.
with self.assertRaises(vfs_plugin.VfsRefreshOperationNotFoundError):
self.handler.Handle(args, token=self.token)
def testHandlerReturnsCorrectStateForFlow(self):
# Create a mock refresh operation.
self.flow_urn = self.CreateRecursiveListFlow(self.client_id, self.token)
args = vfs_plugin.ApiGetVfsRefreshOperationStateArgs(
operation_id=str(self.flow_urn))
# Flow was started and should be running.
result = self.handler.Handle(args, token=self.token)
self.assertEqual(result.state, "RUNNING")
# Terminate flow.
with aff4.FACTORY.Open(
self.flow_urn, aff4_type=flow.GRRFlow, mode="rw",
token=self.token) as flow_obj:
flow_obj.GetRunner().Error("Fake error")
# Recheck status and see if it changed.
result = self.handler.Handle(args, token=self.token)
self.assertEqual(result.state, "FINISHED")
