use of com.alibaba.graphscope.groot.ingestor.BatchSender in project GraphScope by alibaba.
the class BatchSenderTest method testSend.
@Test
void testSend() throws InterruptedException {
Configs configs = Configs.newBuilder().put(CommonConfig.STORE_NODE_COUNT.getKey(), "1").build();
MetaService mockMetaService = mock(MetaService.class);
when(mockMetaService.getPartitionCount()).thenReturn(2);
when(mockMetaService.getStoreIdByPartition(anyInt())).thenReturn(0);
StoreWriter mockStoreWriter = mock(StoreWriter.class);
String requestId = "test_batch_sender";
int queueId = 0;
long snapshotId = 10L;
long offset = 50L;
LabelId labelId = new LabelId(0);
OperationBlob writeVertexBlob1 = new OverwriteVertexOperation(new VertexId(0L), labelId, Collections.EMPTY_MAP).toBlob();
OperationBlob writeVertexBlob2 = new OverwriteVertexOperation(new VertexId(1L), labelId, Collections.EMPTY_MAP).toBlob();
OperationBatch batch = OperationBatch.newBuilder().addOperationBlob(writeVertexBlob1).addOperationBlob(writeVertexBlob2).addOperationBlob(OperationBlob.MARKER_OPERATION_BLOB).build();
CountDownLatch latch = new CountDownLatch(1);
doAnswer(invocationOnMock -> {
StoreDataBatch storeBatch = invocationOnMock.getArgument(1);
CompletionCallback callback = invocationOnMock.getArgument(2);
assertAll(() -> assertEquals(storeBatch.getRequestId(), requestId), () -> assertEquals(storeBatch.getQueueId(), queueId), () -> assertEquals(storeBatch.getSnapshotId(), snapshotId), () -> assertEquals(storeBatch.getOffset(), offset), () -> assertEquals(storeBatch.getDataBatch().size(), 2));
List<Map<Integer, OperationBatch>> dataBatch = storeBatch.getDataBatch();
Map<Integer, OperationBatch> partitionToBatch = dataBatch.get(0);
assertAll(() -> assertEquals(partitionToBatch.get(0).getOperationBlob(0), writeVertexBlob1), () -> assertEquals(partitionToBatch.get(1).getOperationBlob(0), writeVertexBlob2));
assertEquals(dataBatch.get(1).get(-1).getOperationBlob(0), OperationBlob.MARKER_OPERATION_BLOB);
callback.onCompleted(0);
latch.countDown();
return null;
}).when(mockStoreWriter).write(anyInt(), any(), any());
BatchSender batchSender = new BatchSender(configs, mockMetaService, mockStoreWriter, new MetricsCollector(configs));
batchSender.start();
batchSender.asyncSendWithRetry(requestId, queueId, snapshotId, offset, batch);
assertTrue(latch.await(5L, TimeUnit.SECONDS));
batchSender.stop();
}
use of com.alibaba.graphscope.groot.ingestor.BatchSender in project GraphScope by alibaba.
the class IngestProcessorTest method testIngestProcessor.
@Test
void testIngestProcessor() throws IOException {
long tailOffset = 50L;
int queueId = 0;
Configs configs = Configs.newBuilder().build();
LogService mockLogService = mock(LogService.class);
LogReader mockLogReader = mock(LogReader.class);
when(mockLogService.createReader(queueId, tailOffset + 1)).thenReturn(mockLogReader);
LogWriter mockLogWriter = mock(LogWriter.class);
when(mockLogService.createWriter(queueId)).thenReturn(mockLogWriter);
when(mockLogWriter.append(any())).thenReturn(tailOffset + 3);
OperationBatch emptyBatch = OperationBatch.newBuilder().build();
long readSnapshotId = 5L;
ReadLogEntry readLogEntry1 = new ReadLogEntry(tailOffset + 1, readSnapshotId, emptyBatch);
ReadLogEntry readLogEntry2 = new ReadLogEntry(tailOffset + 2, readSnapshotId, emptyBatch);
when(mockLogReader.readNext()).thenReturn(readLogEntry1).thenReturn(readLogEntry2).thenReturn(null);
BatchSender mockBatchSender = mock(BatchSender.class);
AtomicLong ingestSnapshotId = new AtomicLong(10L);
IngestProcessor ingestProcessor = new IngestProcessor(configs, mockLogService, mockBatchSender, queueId, ingestSnapshotId, new MetricsCollector(configs));
ingestProcessor.setTailOffset(tailOffset);
ingestProcessor.start();
verify(mockBatchSender, timeout(5000L)).asyncSendWithRetry(eq(""), eq(queueId), eq(readSnapshotId), eq(tailOffset + 1), any());
verify(mockBatchSender, timeout(5000L)).asyncSendWithRetry(eq(""), eq(queueId), eq(readSnapshotId), eq(tailOffset + 2), any());
verify(mockLogReader).close();
String requestId = "test_ingest_processor";
IngestCallback mockIngestCallback = mock(IngestCallback.class);
ingestProcessor.ingestBatch(requestId, emptyBatch, mockIngestCallback);
verify(mockBatchSender, timeout(5000L)).asyncSendWithRetry(requestId, queueId, ingestSnapshotId.get(), tailOffset + 3, emptyBatch);
verify(mockIngestCallback, timeout(5000L)).onSuccess(ingestSnapshotId.get());
ingestProcessor.stop();
verify(mockLogWriter, timeout(5000L)).close();
}
Aggregations