def MockMethodIsCalled():
try:
# Mock should be called twice: once for HEAD (to check permissions)
# and once for GET methods.
mock_method.assert_called_with(
api_hunt.ApiGetExportedHuntResultsArgs(
hunt_id=hunt_urn.Basename(), plugin_name=plugin),
token=mock.ANY)
return True
except AssertionError:
return False
def testWorksCorrectlyWithTestOutputPluginOnFlowWithSingleResult(self):
result = self.handler.Handle(hunt_plugin.ApiGetExportedHuntResultsArgs(
hunt_id=self.hunt.urn.Basename(),
plugin_name=test_plugins.TestInstantOutputPlugin.plugin_name),
token=self.token)
chunks = list(result.GenerateContent())
self.assertListEqual(chunks, [
"Start: %s" % utils.SmartStr(self.hunt.urn),
"Values of type: RDFString",
"First pass: oh (source=%s)" % self.client_ids[0],
"First pass: oh (source=%s)" % self.client_ids[1],
"First pass: oh (source=%s)" % self.client_ids[2],
"First pass: oh (source=%s)" % self.client_ids[3],
"First pass: oh (source=%s)" % self.client_ids[4],
"Second pass: oh (source=%s)" % self.client_ids[0],
"Second pass: oh (source=%s)" % self.client_ids[1],
"Second pass: oh (source=%s)" % self.client_ids[2],
"Second pass: oh (source=%s)" % self.client_ids[3],
"Second pass: oh (source=%s)" % self.client_ids[4],
"Finish: %s" % utils.SmartStr(self.hunt.urn)
]) # pyformat: disable
def testWorksCorrectlyWithTestOutputPluginOnFlowWithSingleResult(self):
result = self.handler.Handle(hunt_plugin.ApiGetExportedHuntResultsArgs(
hunt_id=self.hunt_id,
plugin_name=test_plugins.TestInstantOutputPlugin.plugin_name),
token=self.token)
chunks = list(result.GenerateContent())
self.assertListEqual(chunks, [
"Start: aff4:/hunts/%s" % self.hunt_id,
"Values of type: RDFString",
"First pass: oh (source=aff4:/%s)" % self.client_ids[0],
"First pass: oh (source=aff4:/%s)" % self.client_ids[1],
"First pass: oh (source=aff4:/%s)" % self.client_ids[2],
"First pass: oh (source=aff4:/%s)" % self.client_ids[3],
"First pass: oh (source=aff4:/%s)" % self.client_ids[4],
"Second pass: oh (source=aff4:/%s)" % self.client_ids[0],
"Second pass: oh (source=aff4:/%s)" % self.client_ids[1],
"Second pass: oh (source=aff4:/%s)" % self.client_ids[2],
"Second pass: oh (source=aff4:/%s)" % self.client_ids[3],
"Second pass: oh (source=aff4:/%s)" % self.client_ids[4],
"Finish: aff4:/hunts/%s" % self.hunt_id,
])
