def ProcessResponses(self, plugin_args=None, responses=None, process_responses_separately=False): plugin = bigquery_plugin.BigQueryOutputPlugin( source_urn=self.results_urn, output_base_urn=self.base_urn, args=plugin_args, token=self.token) plugin.Initialize() messages = [] for response in responses: messages.append( rdf_flows.GrrMessage(source=self.client_id, payload=response)) with test_lib.FakeTime(1445995873): with mock.patch.object(bigquery, "GetBigQueryClient") as mock_bigquery: if process_responses_separately: for message in messages: plugin.ProcessResponses([message]) else: plugin.ProcessResponses(messages) plugin.Flush() return [ x[0] for x in mock_bigquery.return_value.InsertData.call_args_list ]
def testBigQueryPluginFallbackToAFF4(self): plugin_args = bigquery_plugin.BigQueryOutputPluginArgs() responses = [ rdf_client.StatEntry( aff4path=self.client_id.Add("/fs/os/中国新闻网新闻中"), pathspec=rdf_paths.PathSpec(path="/中国新闻网新闻中")), rdf_client.Process(pid=42), rdf_client.Process(pid=43), rdf_client.SoftwarePackage(name="test.deb") ] plugin = bigquery_plugin.BigQueryOutputPlugin( source_urn=self.results_urn, output_base_urn=self.base_urn, args=plugin_args, token=self.token) plugin.Initialize() messages = [] for response in responses: messages.append( rdf_flows.GrrMessage(source=self.client_id, payload=response)) with test_lib.FakeTime(1445995873): with mock.patch.object(bigquery, "GetBigQueryClient") as mock_bigquery: mock_bigquery.return_value.configure_mock(**{ "InsertData.side_effect": bigquery.BigQueryJobUploadError() }) with test_lib.ConfigOverrider( {"BigQuery.max_upload_failures": 2}): for message in messages: plugin.ProcessResponses([message]) plugin.Flush() # We have 3 output types but a limit of 2 upload failures, so we # shouldn't try the third one. self.assertEqual( mock_bigquery.return_value.InsertData.call_count, 2) # We should have written a data file and a schema file for each type. for output_name in [ "ExportedFile", "ExportedProcess", "AutoExportedSoftwarePackage" ]: schema_fd = aff4.FACTORY.Open(self.base_urn.Add( "C-1000000000000000_Results_%s_1445995873.schema" % output_name), token=self.token) data_fd = aff4.FACTORY.Open(self.base_urn.Add( "C-1000000000000000_Results_%s_1445995873.data" % output_name), token=self.token) actual_fd = gzip.GzipFile(None, "r", 9, data_fd) if output_name == "ExportedFile": self.CompareSchemaToKnownGood(json.load(schema_fd)) self.assertEqual( json.load(actual_fd)["urn"], self.client_id.Add("/fs/os/中国新闻网新闻中")) elif output_name == "ExportedProcess": self.assertEqual(json.load(schema_fd)[1]["name"], "pid") expected_pids = ["42", "43"] for i, line in enumerate(actual_fd): self.assertEqual(json.loads(line)["pid"], expected_pids[i]) else: self.assertEqual(json.load(schema_fd)[1]["name"], "name") self.assertEqual(json.load(actual_fd)["name"], "test.deb") # Process the same messages to make sure we're re-using the filehandles. with test_lib.FakeTime(1445995878): with mock.patch.object(bigquery, "GetBigQueryClient") as mock_bigquery: mock_bigquery.return_value.configure_mock(**{ "InsertData.side_effect": bigquery.BigQueryJobUploadError() }) with test_lib.ConfigOverrider( {"BigQuery.max_upload_failures": 2}): for message in messages: plugin.ProcessResponses([message]) plugin.Flush() # We shouldn't call insertdata at all because we have passed max # failures already self.assertEqual( mock_bigquery.return_value.InsertData.call_count, 0) expected_line_counts = { "ExportedFile": 2, "ExportedProcess": 4, "AutoExportedSoftwarePackage": 2 } for output_name in [ "ExportedFile", "ExportedProcess", "AutoExportedSoftwarePackage" ]: data_fd = aff4.FACTORY.Open(self.base_urn.Add( "C-1000000000000000_Results_%s_1445995873.data" % output_name), token=self.token) actual_fd = gzip.GzipFile(None, "r", 9, data_fd) self.assertEqual(sum(1 for line in actual_fd), expected_line_counts[output_name])