Examples with BulkMessage org.apache.metron.common.writer.BulkMessage used on opensource projects

Example 21 with BulkMessage

use of org.apache.metron.common.writer.BulkMessage in project metron by apache.

the class HdfsWriterTest method testWriteNoOutputFunction.

@Test
@SuppressWarnings("unchecked")
public void testWriteNoOutputFunction() throws Exception {
    FileNameFormat format = new DefaultFileNameFormat().withPath(folder.toString()).withExtension(".json").withPrefix("prefix-");
    HdfsWriter writer = new HdfsWriter().withFileNameFormat(format);
    IndexingConfigurations indexingConfig = new IndexingConfigurations();
    WriterConfiguration config = new IndexingWriterConfiguration(WRITER_NAME, indexingConfig);
    writer.init(new HashMap<String, String>(), config);
    writer.initFileNameFormat(createTopologyContext());
    JSONObject message = new JSONObject();
    message.put("test.key", "test.value");
    message.put("test.key2", "test.value2");
    JSONObject message2 = new JSONObject();
    message2.put("test.key", "test.value3");
    message2.put("test.key2", "test.value2");
    List<BulkMessage<JSONObject>> messages = new ArrayList<BulkMessage<JSONObject>>() {

        {
            add(new BulkMessage("message1", message));
            add(new BulkMessage("message2", message2));
        }
    };
    writer.write(SENSOR_NAME, config, messages);
    writer.close();
    ArrayList<String> expected = new ArrayList<>();
    expected.add(message.toJSONString());
    expected.add(message2.toJSONString());
    Collections.sort(expected);
    // Default to just putting it in the base folder + the sensor name
    File outputFolder = new File(folder.getAbsolutePath() + "/" + SENSOR_NAME);
    assertTrue(outputFolder.exists() && outputFolder.isDirectory());
    assertEquals(1, outputFolder.listFiles().length);
    for (File file : outputFolder.listFiles()) {
        List<String> lines = Files.readAllLines(file.toPath());
        Collections.sort(lines);
        assertEquals(expected, lines);
    }
}

Example 22 with BulkMessage

use of org.apache.metron.common.writer.BulkMessage in project metron by apache.

the class KafkaWriterTest method testWriterShouldReturnResponse.

@Test
public void testWriterShouldReturnResponse() throws Exception {
    KafkaWriter writer = spy(new KafkaWriter());
    writer.setKafkaProducer(kafkaProducer);
    List<BulkMessage<JSONObject>> messages = new ArrayList<>();
    JSONObject successMessage = new JSONObject();
    successMessage.put("value", "success");
    JSONObject errorMessage = new JSONObject();
    errorMessage.put("value", "error");
    JSONObject droppedMessage = new JSONObject();
    droppedMessage.put("value", "dropped");
    messages.add(new BulkMessage<>("successId", successMessage));
    messages.add(new BulkMessage<>("errorId", errorMessage));
    messages.add(new BulkMessage<>("droppedId", droppedMessage));
    doReturn(Optional.of("successTopic")).when(writer).getKafkaTopic(successMessage);
    doReturn(Optional.of("errorTopic")).when(writer).getKafkaTopic(errorMessage);
    doReturn(Optional.empty()).when(writer).getKafkaTopic(droppedMessage);
    Future successFuture = mock(Future.class);
    Future errorFuture = mock(Future.class);
    ExecutionException throwable = new ExecutionException(new Exception("kafka error"));
    when(kafkaProducer.send(new ProducerRecord<String, String>("errorTopic", "{\"value\":\"error\"}"))).thenReturn(errorFuture);
    when(kafkaProducer.send(new ProducerRecord<String, String>("successTopic", "{\"value\":\"success\"}"))).thenReturn(successFuture);
    when(errorFuture.get()).thenThrow(throwable);
    BulkWriterResponse response = new BulkWriterResponse();
    response.addSuccess(new MessageId("successId"));
    response.addError(throwable, new MessageId("errorId"));
    assertEquals(response, writer.write(SENSOR_TYPE, createConfiguration(new HashMap<>()), messages));
    verify(kafkaProducer, times(1)).flush();
    verify(kafkaProducer, times(1)).send(new ProducerRecord<String, String>("successTopic", "{\"value\":\"success\"}"));
    verify(kafkaProducer, times(1)).send(new ProducerRecord<String, String>("errorTopic", "{\"value\":\"error\"}"));
    verifyNoMoreInteractions(kafkaProducer);
}

Example 23 with BulkMessage

use of org.apache.metron.common.writer.BulkMessage in project metron by apache.

the class BulkWriterComponentTest method writeShouldProperlyAckTuplesInBatch.

@Test
public void writeShouldProperlyAckTuplesInBatch() throws Exception {
    BulkWriterComponent<JSONObject> bulkWriterComponent = new BulkWriterComponent<>(Collections.singletonList(flushPolicy));
    BulkWriterResponse response = new BulkWriterResponse();
    response.addAllSuccesses(messageIds);
    when(bulkMessageWriter.write(sensorType, configurations, messages)).thenReturn(response);
    bulkWriterComponent.write(sensorType, messages.get(0), bulkMessageWriter, configurations);
    verify(bulkMessageWriter, times(0)).write(eq(sensorType), eq(configurations), any());
    verify(flushPolicy, times(1)).shouldFlush(sensorType, configurations, messages.subList(0, 1));
    verify(flushPolicy, times(0)).onFlush(any(), any());
    reset(flushPolicy);
    when(flushPolicy.shouldFlush(sensorType, configurations, messages)).thenReturn(true);
    bulkWriterComponent.write(sensorType, messages.get(1), bulkMessageWriter, configurations);
    BulkWriterResponse expectedResponse = new BulkWriterResponse();
    expectedResponse.addAllSuccesses(messageIds);
    verify(bulkMessageWriter, times(1)).write(sensorType, configurations, Arrays.asList(new BulkMessage<>(messageId1, message1), new BulkMessage<>(messageId2, message2)));
    verify(flushPolicy, times(1)).shouldFlush(sensorType, configurations, messages);
    verify(flushPolicy, times(1)).onFlush(sensorType, expectedResponse);
    verifyNoMoreInteractions(bulkMessageWriter, flushPolicy);
}

Example 24 with BulkMessage

use of org.apache.metron.common.writer.BulkMessage in project metron by apache.

the class SimpleHBaseEnrichmentWriterTest method testFilteredKeys.

@Test
public void testFilteredKeys() throws Exception {
    final String sensorType = "dummy";
    SimpleHbaseEnrichmentWriter writer = new SimpleHbaseEnrichmentWriter();
    WriterConfiguration configuration = createConfig(1, new HashMap<String, Object>(BASE_WRITER_CONFIG) {

        {
            put(SimpleHbaseEnrichmentWriter.Configurations.KEY_COLUMNS.getKey(), "ip");
            put(SimpleHbaseEnrichmentWriter.Configurations.VALUE_COLUMNS.getKey(), ImmutableList.of("user", "ip"));
        }
    });
    writer.configure(sensorType, configuration);
    writer.write(SENSOR_TYPE, configuration, new ArrayList<BulkMessage<JSONObject>>() {

        {
            add(new BulkMessage<>("messageId", new JSONObject(ImmutableMap.of("ip", "localhost", "user", "cstella", "foo", "bar"))));
        }
    });
    List<LookupKV<EnrichmentKey, EnrichmentValue>> values = getValues();
    assertEquals(1, values.size());
    assertEquals("localhost", values.get(0).getKey().indicator);
    assertEquals("cstella", values.get(0).getValue().getMetadata().get("user"));
    assertEquals("localhost", values.get(0).getValue().getMetadata().get("ip"));
    assertNull(values.get(0).getValue().getMetadata().get("foo"));
    assertEquals(2, values.get(0).getValue().getMetadata().size());
}

Example 25 with BulkMessage

use of org.apache.metron.common.writer.BulkMessage in project metron by apache.

the class SolrWriterTest method testWriter.

@Test
public void testWriter() throws Exception {
    IndexingConfigurations configurations = SampleUtil.getSampleIndexingConfigs();
    JSONObject message1 = new JSONObject();
    message1.put(Constants.GUID, "guid-1");
    message1.put(Constants.SENSOR_TYPE, "test");
    message1.put("intField", 100);
    message1.put("doubleField", 100.0);
    JSONObject message2 = new JSONObject();
    message2.put(Constants.GUID, "guid-2");
    message2.put(Constants.SENSOR_TYPE, "test");
    message2.put("intField", 200);
    message2.put("doubleField", 200.0);
    List<BulkMessage<JSONObject>> messages = new ArrayList<>();
    messages.add(new BulkMessage<>("message1", message1));
    messages.add(new BulkMessage<>("message2", message2));
    String collection = "metron";
    MetronSolrClient solr = Mockito.mock(MetronSolrClient.class);
    SolrWriter writer = new SolrWriter().withMetronSolrClient(solr);
    writer.init(null, new IndexingWriterConfiguration("solr", configurations));
    verify(solr, times(1)).setDefaultCollection(collection);
    collection = "metron2";
    Map<String, Object> globalConfig = configurations.getGlobalConfig();
    globalConfig.put("solr.collection", collection);
    configurations.updateGlobalConfig(globalConfig);
    writer = new SolrWriter().withMetronSolrClient(solr);
    writer.init(null, new IndexingWriterConfiguration("solr", configurations));
    verify(solr, times(1)).setDefaultCollection(collection);
    writer.write("test", new IndexingWriterConfiguration("solr", configurations), messages);
    verify(solr, times(1)).add(eq("yaf"), argThat(new SolrInputDocumentMatcher(ImmutableList.of(message1, message2))));
    verify(solr, times(1)).commit("yaf", (boolean) SolrWriter.SolrProperties.COMMIT_WAIT_FLUSH.defaultValue.get(), (boolean) SolrWriter.SolrProperties.COMMIT_WAIT_SEARCHER.defaultValue.get(), (boolean) SolrWriter.SolrProperties.COMMIT_SOFT.defaultValue.get());
}