Examples with RequestCase org.apache.bookkeeper.stream.proto.storage.StorageContainerRequest.RequestCase used on opensource projects

Example 1 with RequestCase

use of org.apache.bookkeeper.stream.proto.storage.StorageContainerRequest.RequestCase 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());
}