use of org.apache.bookkeeper.stream.proto.storage.TableServiceGrpc.TableServiceImplBase in project bookkeeper by apache.
the class TableRequestProcessorTest method testProcess.
private void testProcess(RequestCase type) throws Exception {
@Cleanup("shutdown") ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
StorageContainerChannel scChannel = mock(StorageContainerChannel.class);
CompletableFuture<StorageServerChannel> serverChannelFuture = FutureUtils.createFuture();
when(scChannel.getStorageContainerChannelFuture()).thenReturn(serverChannelFuture);
StorageContainerResponse.Builder respBuilder = StorageContainerResponse.newBuilder().setCode(StatusCode.SUCCESS);
switch(type) {
case KV_PUT_REQ:
respBuilder.setKvPutResp(PutResponse.newBuilder().build());
break;
case KV_DELETE_REQ:
respBuilder.setKvDeleteResp(DeleteRangeResponse.newBuilder().build());
break;
case KV_RANGE_REQ:
respBuilder.setKvRangeResp(RangeResponse.newBuilder().build());
break;
case KV_INCR_REQ:
respBuilder.setKvIncrResp(IncrementResponse.newBuilder().build());
break;
case KV_TXN_REQ:
respBuilder.setKvTxnResp(TxnResponse.newBuilder().build());
break;
default:
break;
}
StorageContainerResponse response = respBuilder.build();
AtomicReference<StorageContainerRequest> receivedRequest = new AtomicReference<>(null);
AtomicReference<RequestCase> receivedRequestType = new AtomicReference<>(null);
TableServiceImplBase tableService = new TableServiceImplBase() {
@Override
public void range(StorageContainerRequest request, StreamObserver<StorageContainerResponse> responseObserver) {
receivedRequest.set(request);
receivedRequestType.set(KV_RANGE_REQ);
complete(responseObserver);
}
@Override
public void put(StorageContainerRequest request, StreamObserver<StorageContainerResponse> responseObserver) {
receivedRequest.set(request);
receivedRequestType.set(KV_PUT_REQ);
complete(responseObserver);
}
@Override
public void delete(StorageContainerRequest request, StreamObserver<StorageContainerResponse> responseObserver) {
receivedRequest.set(request);
receivedRequestType.set(KV_DELETE_REQ);
complete(responseObserver);
}
@Override
public void txn(StorageContainerRequest request, StreamObserver<StorageContainerResponse> responseObserver) {
receivedRequest.set(request);
receivedRequestType.set(KV_TXN_REQ);
complete(responseObserver);
}
@Override
public void increment(StorageContainerRequest request, StreamObserver<StorageContainerResponse> responseObserver) {
receivedRequest.set(request);
receivedRequestType.set(KV_INCR_REQ);
complete(responseObserver);
}
private void complete(StreamObserver<StorageContainerResponse> responseStreamObserver) {
responseStreamObserver.onNext(response);
responseStreamObserver.onCompleted();
}
};
serviceRegistry.addService(tableService.bindService());
StorageServerChannel ssChannel = new StorageServerChannel(InProcessChannelBuilder.forName(serverName).directExecutor().build(), Optional.empty());
serverChannelFuture.complete(ssChannel);
StorageContainerRequest.Builder requestBuilder = StorageContainerRequest.newBuilder();
switch(type) {
case KV_PUT_REQ:
requestBuilder.setKvPutReq(PutRequest.newBuilder().build());
break;
case KV_DELETE_REQ:
requestBuilder.setKvDeleteReq(DeleteRangeRequest.newBuilder().build());
break;
case KV_RANGE_REQ:
requestBuilder.setKvRangeReq(RangeRequest.newBuilder().build());
break;
case KV_INCR_REQ:
requestBuilder.setKvIncrReq(IncrementRequest.newBuilder().build());
break;
case KV_TXN_REQ:
requestBuilder.setKvTxnReq(TxnRequest.newBuilder().build());
break;
default:
break;
}
StorageContainerRequest request = requestBuilder.build();
TableRequestProcessor<String> processor = TableRequestProcessor.of(request, resp -> "test", scChannel, scheduler);
assertEquals("test", FutureUtils.result(processor.process()));
assertSame(request, receivedRequest.get());
assertEquals(type, receivedRequestType.get());
}
Aggregations