Search in sources :

Example 1 with SubscriberImplBase

use of com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase in project beam by apache.

the class PubsubGrpcClientTest method pullOneMessage.

@Test
public void pullOneMessage() throws IOException {
    initializeClient(null, null);
    String expectedSubscription = SUBSCRIPTION.getPath();
    final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
    Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
    PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).build();
    ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
    final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
    final List<PullRequest> requestsReceived = new ArrayList<>();
    SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {

        @Override
        public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
            requestsReceived.add(request);
            responseObserver.onNext(response);
            responseObserver.onCompleted();
        }
    };
    Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
    try {
        List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
        assertEquals(1, actualMessages.size());
        IncomingMessage actualMessage = actualMessages.get(0);
        assertEquals(ACK_ID, actualMessage.ackId());
        assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
        assertEquals(MESSAGE_ID, actualMessage.recordId());
        assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
        assertEquals(PUB_TIME_MS, actualMessage.timestampMsSinceEpoch());
        assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
    } finally {
        server.shutdownNow();
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) SubscriberImplBase(com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase) Server(io.grpc.Server) PullRequest(com.google.pubsub.v1.PullRequest) ArrayList(java.util.ArrayList) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) Timestamp(com.google.protobuf.Timestamp) PubsubMessage(com.google.pubsub.v1.PubsubMessage) PullResponse(com.google.pubsub.v1.PullResponse) IncomingMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage) Test(org.junit.Test)

Example 2 with SubscriberImplBase

use of com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase in project beam by apache.

the class PubsubGrpcClientTest method pullOneMessageUsingAttributes.

@Test
public void pullOneMessageUsingAttributes() throws IOException {
    initializeClient(TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE);
    String expectedSubscription = SUBSCRIPTION.getPath();
    final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
    Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
    PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).putAllAttributes(ATTRIBUTES).putAllAttributes(ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME_MS), ID_ATTRIBUTE, RECORD_ID)).build();
    ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
    final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
    final List<PullRequest> requestsReceived = new ArrayList<>();
    SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {

        @Override
        public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
            requestsReceived.add(request);
            responseObserver.onNext(response);
            responseObserver.onCompleted();
        }
    };
    Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
    try {
        List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
        assertEquals(1, actualMessages.size());
        IncomingMessage actualMessage = actualMessages.get(0);
        assertEquals(ACK_ID, actualMessage.ackId());
        assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
        assertEquals(RECORD_ID, actualMessage.recordId());
        assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
        assertEquals(MESSAGE_TIME_MS, actualMessage.timestampMsSinceEpoch());
        assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
    } finally {
        server.shutdownNow();
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) SubscriberImplBase(com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase) Server(io.grpc.Server) PullRequest(com.google.pubsub.v1.PullRequest) ArrayList(java.util.ArrayList) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) Timestamp(com.google.protobuf.Timestamp) PubsubMessage(com.google.pubsub.v1.PubsubMessage) PullResponse(com.google.pubsub.v1.PullResponse) IncomingMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage) Test(org.junit.Test)

Aggregations

ByteString (com.google.protobuf.ByteString)2 Timestamp (com.google.protobuf.Timestamp)2 PubsubMessage (com.google.pubsub.v1.PubsubMessage)2 PullRequest (com.google.pubsub.v1.PullRequest)2 PullResponse (com.google.pubsub.v1.PullResponse)2 ReceivedMessage (com.google.pubsub.v1.ReceivedMessage)2 SubscriberImplBase (com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase)2 Server (io.grpc.Server)2 StreamObserver (io.grpc.stub.StreamObserver)2 ArrayList (java.util.ArrayList)2 IncomingMessage (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage)2 Test (org.junit.Test)2