Examples with Record com.amazon.dataprepper.model.record.Record used on opensource projects

Example 31 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class BufferAccumulatorTest method flush_after_add_writes_to_buffer.

@Test
void flush_after_add_writes_to_buffer() throws Exception {
    final BufferAccumulator objectUnderTest = createObjectUnderTest();
    final Record record = createRecord();
    objectUnderTest.add(record);
    verifyNoInteractions(buffer);
    final List<Record<?>> actualRecordsWritten = new ArrayList<>();
    doAnswer(a -> actualRecordsWritten.addAll(a.getArgument(0, Collection.class))).when(buffer).writeAll(anyCollection(), anyInt());
    objectUnderTest.flush();
    verify(buffer).writeAll(anyCollection(), eq(timeoutMillis));
    assertThat(actualRecordsWritten.size(), equalTo(1));
    assertThat(actualRecordsWritten, equalTo(Collections.singletonList(record)));
}

Example 32 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class S3ObjectWorkerTest method parseS3Object_calls_Codec_parse_with_Consumer_that_adds_to_BufferAccumulator.

@Test
void parseS3Object_calls_Codec_parse_with_Consumer_that_adds_to_BufferAccumulator() throws Exception {
    final ResponseInputStream<GetObjectResponse> objectInputStream = mock(ResponseInputStream.class);
    when(s3Client.getObject(any(GetObjectRequest.class))).thenReturn(objectInputStream);
    final BufferAccumulator bufferAccumulator = mock(BufferAccumulator.class);
    try (final MockedStatic<BufferAccumulator> bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) {
        bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, recordsToAccumulate, bufferTimeout)).thenReturn(bufferAccumulator);
        createObjectUnderTest().parseS3Object(s3ObjectReference);
    }
    final ArgumentCaptor<Consumer<Record<Event>>> eventConsumerArgumentCaptor = ArgumentCaptor.forClass(Consumer.class);
    verify(codec).parse(any(InputStream.class), eventConsumerArgumentCaptor.capture());
    final Consumer<Record<Event>> consumerUnderTest = eventConsumerArgumentCaptor.getValue();
    final Record record = mock(Record.class);
    consumerUnderTest.accept(record);
    verify(bufferAccumulator).add(record);
}

Example 33 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class NewlineDelimitedCodecTest method parse_calls_Consumer_for_each_line.

@ParameterizedTest
@ValueSource(ints = { 0, 1, 2, 10, 50 })
void parse_calls_Consumer_for_each_line(int numberOfLines) throws IOException {
    final List<String> linesList = generateLinesAsList(numberOfLines);
    final InputStream inputStream = createInputStream(linesList);
    List<Record<Event>> actualEvents = new ArrayList<>();
    createObjectUnderTest().parse(inputStream, actualEvents::add);
    assertThat(actualEvents.size(), equalTo(numberOfLines));
    for (int i = 0; i < actualEvents.size(); i++) {
        final Record<Event> record = actualEvents.get(i);
        assertThat(record, notNullValue());
        assertThat(record.getData(), notNullValue());
        assertThat(record.getData().get("message", String.class), equalTo(linesList.get(i)));
    }
}

Example 34 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class NewlineDelimitedCodecTest method parse_calls_Consumer_for_each_line_after_skipping.

@ParameterizedTest
@ValueSource(ints = { 1, 2, 10, 50 })
void parse_calls_Consumer_for_each_line_after_skipping(int numberOfLines) throws IOException {
    final List<String> linesList = generateLinesAsList(numberOfLines);
    final InputStream inputStream = createInputStream(linesList);
    final int skipLines = 1;
    when(config.getSkipLines()).thenReturn(skipLines);
    final List<Record<Event>> actualEvents = new ArrayList<>();
    createObjectUnderTest().parse(inputStream, actualEvents::add);
    assertThat(actualEvents.size(), equalTo(numberOfLines - skipLines));
    for (int i = 0; i < actualEvents.size(); i++) {
        final Record<Event> record = actualEvents.get(i);
        assertThat(record, notNullValue());
        assertThat(record.getData(), notNullValue());
        assertThat(record.getData().get("message", String.class), equalTo(linesList.get(i + skipLines)));
    }
}

Example 35 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class ServiceMapStatefulPrepperTest method testPrepareForShutdownWithEventRecordData.

@Test
public void testPrepareForShutdownWithEventRecordData() {
    final File path = new File(ServiceMapPrepperConfig.DEFAULT_DB_PATH);
    final ServiceMapStatefulPrepper serviceMapStateful = new ServiceMapStatefulPrepper(100, path, Clock.systemUTC(), 1, PLUGIN_SETTING);
    final byte[] rootSpanId1Bytes = ServiceMapTestUtils.getRandomBytes(8);
    final byte[] traceId1Bytes = ServiceMapTestUtils.getRandomBytes(16);
    final String rootSpanId1 = Hex.encodeHexString(rootSpanId1Bytes);
    final String traceId1 = Hex.encodeHexString(traceId1Bytes);
    final String traceGroup1 = "reset_password";
    final Span frontendSpans1 = ServiceMapTestUtils.getSpan(FRONTEND_SERVICE, traceGroup1, rootSpanId1, "", traceId1, io.opentelemetry.proto.trace.v1.Span.SpanKind.SPAN_KIND_CLIENT);
    final Span authenticationSpansServer = ServiceMapTestUtils.getSpan(AUTHENTICATION_SERVICE, "reset", Hex.encodeHexString(ServiceMapTestUtils.getRandomBytes(8)), frontendSpans1.getSpanId(), traceId1, io.opentelemetry.proto.trace.v1.Span.SpanKind.SPAN_KIND_SERVER);
    serviceMapStateful.execute(Arrays.asList(new Record<>(frontendSpans1), new Record<>(authenticationSpansServer)));
    assertFalse(serviceMapStateful.isReadyForShutdown());
    serviceMapStateful.prepareForShutdown();
    serviceMapStateful.execute(Collections.emptyList());
    assertTrue(serviceMapStateful.isReadyForShutdown());
    serviceMapStateful.shutdown();
}