use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.
the class PipelineTests method testGetSinks.
@Test
public void testGetSinks() {
final Source<Record<String>> testSource = new TestSource();
final TestSink testSink = new TestSink();
final Pipeline testPipeline = new Pipeline(TEST_PIPELINE_NAME, testSource, new BlockingBuffer(TEST_PIPELINE_NAME), Collections.emptyList(), Collections.singletonList(testSink), TEST_PROCESSOR_THREADS, TEST_READ_BATCH_TIMEOUT);
assertEquals(1, testPipeline.getSinks().size());
assertEquals(testSink, testPipeline.getSinks().iterator().next());
}
use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.
the class PeerForwarderTest method testSingleRemoteIpForwardRequestEncodeError.
@Test
public void testSingleRemoteIpForwardRequestEncodeError() throws NoSuchFieldException, IllegalAccessException, DecoderException, UnsupportedEncodingException {
final List<String> testIps = generateTestIps(2);
final String peerIp = testIps.get(1);
when(peerClientPool.getClient(peerIp)).thenReturn(client);
MetricsTestUtil.initMetrics();
final PeerForwarder testPeerForwarder = generatePeerForwarder(testIps, 3);
when(oTelProtoEncoder.convertToResourceSpans(any(Span.class))).thenThrow(new DecoderException());
reflectivelySetEncoder(testPeerForwarder, oTelProtoEncoder);
final List<Record<Object>> exportedRecords = testPeerForwarder.doExecute(TEST_SPANS_B.stream().map(span -> new Record<Object>(span)).collect(Collectors.toList()));
verifyNoInteractions(client);
Assert.assertEquals(3, exportedRecords.size());
final List<Span> exportedSpans = exportedRecords.stream().map(record -> (Span) record.getData()).collect(Collectors.toList());
assertTrue(exportedSpans.containsAll(TEST_SPANS_B));
assertTrue(TEST_SPANS_B.containsAll(exportedSpans));
}
use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.
the class PeerForwarderTest method testSingleRemoteIpForwardRequestFutureError.
@Test
public void testSingleRemoteIpForwardRequestFutureError() throws ExecutionException, InterruptedException {
try (final MockedStatic<CompletableFuture> completableFutureMockedStatic = mockStatic(CompletableFuture.class)) {
completableFutureMockedStatic.when(() -> CompletableFuture.supplyAsync(ArgumentMatchers.<Supplier<ExportTraceServiceRequest>>any(), any(ExecutorService.class))).thenReturn(completableFuture);
when(completableFuture.get()).thenThrow(new InterruptedException());
final List<String> testIps = generateTestIps(2);
final Channel channel = mock(Channel.class);
final String peerIp = testIps.get(1);
final String fullPeerIp = String.format("%s:21890", peerIp);
when(channel.authority()).thenReturn(fullPeerIp);
when(peerClientPool.getClient(peerIp)).thenReturn(client);
when(client.getChannel()).thenReturn(channel);
MetricsTestUtil.initMetrics();
final PeerForwarder testPeerForwarder = generatePeerForwarder(testIps, 3);
final List<Record<Object>> exportedRecords = testPeerForwarder.doExecute(TEST_SPANS_B.stream().map(span -> new Record<Object>(span)).collect(Collectors.toList()));
verify(completableFuture, times(1)).get();
Assert.assertEquals(3, exportedRecords.size());
final List<Span> exportedSpans = exportedRecords.stream().map(record -> (Span) record.getData()).collect(Collectors.toList());
assertTrue(exportedSpans.containsAll(TEST_SPANS_B));
assertTrue(TEST_SPANS_B.containsAll(exportedSpans));
}
}
use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.
the class PeerForwarderTest method testSingleRemoteIpBothLocalAndForwardedRequestWithEventRecordData.
@Test
public void testSingleRemoteIpBothLocalAndForwardedRequestWithEventRecordData() throws DecoderException {
final List<String> testIps = generateTestIps(2);
final Channel channel = mock(Channel.class);
final String peerIp = testIps.get(1);
when(channel.authority()).thenReturn(String.format("%s:21890", peerIp));
when(peerClientPool.getClient(peerIp)).thenReturn(client);
when(client.getChannel()).thenReturn(channel);
final Map<String, List<ExportTraceServiceRequest>> requestsByIp = testIps.stream().collect(Collectors.toMap(ip -> ip, ip -> new ArrayList<>()));
doAnswer(invocation -> {
final ExportTraceServiceRequest exportTraceServiceRequest = invocation.getArgument(0);
requestsByIp.get(peerIp).add(exportTraceServiceRequest);
return null;
}).when(client).export(any(ExportTraceServiceRequest.class));
MetricsTestUtil.initMetrics();
final PeerForwarder testPeerForwarder = generatePeerForwarder(testIps, 3);
final List<Record<Object>> exportedRecords = testPeerForwarder.doExecute(TEST_SPANS_ALL.stream().map(span -> new Record<Object>(span)).collect(Collectors.toList()));
final List<Span> expectedLocalSpans = Arrays.asList(SPAN_1, SPAN_2, SPAN_3);
Assert.assertEquals(3, exportedRecords.size());
final List<Span> localSpans = exportedRecords.stream().map(record -> (Span) record.getData()).collect(Collectors.toList());
assertTrue(localSpans.containsAll(expectedLocalSpans));
assertTrue(expectedLocalSpans.containsAll(localSpans));
Assert.assertEquals(1, requestsByIp.get(peerIp).size());
final ExportTraceServiceRequest forwardedRequest = requestsByIp.get(peerIp).get(0);
final List<ResourceSpans> forwardedResourceSpans = forwardedRequest.getResourceSpansList();
assertEquals(3, forwardedResourceSpans.size());
forwardedResourceSpans.forEach(rs -> {
assertEquals(TEST_SERVICE_B, extractServiceName(rs));
assertEquals(1, rs.getInstrumentationLibrarySpansCount());
final InstrumentationLibrarySpans ils = rs.getInstrumentationLibrarySpans(0);
assertEquals(1, ils.getSpansCount());
final io.opentelemetry.proto.trace.v1.Span sp = ils.getSpans(0);
assertEquals(TEST_TRACE_ID_2, Hex.encodeHexString(sp.getTraceId().toByteArray()));
});
}
use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.
the class BufferAccumulatorTest method adding_to_accumulation_count_should_write_to_Buffer.
@ParameterizedTest
@ValueSource(ints = { 2, 10, 20 })
void adding_to_accumulation_count_should_write_to_Buffer(final int accumulationCount) throws Exception {
recordsToAccumulate = accumulationCount;
final BufferAccumulator objectUnderTest = createObjectUnderTest();
final List<Record<?>> knownRecords = new ArrayList<>();
for (int i = 0; i < accumulationCount - 1; i++) {
final Record record = createRecord();
knownRecords.add(record);
objectUnderTest.add(record);
}
verifyNoInteractions(buffer);
final Collection<Record<?>> actualRecordsWritten = new ArrayList<>();
doAnswer(a -> actualRecordsWritten.addAll(a.getArgument(0, Collection.class))).when(buffer).writeAll(anyCollection(), anyInt());
final Record record = createRecord();
knownRecords.add(record);
objectUnderTest.add(record);
verify(buffer).writeAll(anyCollection(), eq(timeoutMillis));
assertThat(actualRecordsWritten.size(), equalTo(accumulationCount));
assertThat(actualRecordsWritten, equalTo(knownRecords));
}
Aggregations