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, ])