com.google.pubsub.v1.AcknowledgeRequest Java Examples

@Override
public void acknowledge(List<String> acknowledgementIds) {
	if (acknowledgementIds.isEmpty()) {
		return;
	}

	//grpc servers won't accept acknowledge requests that are too large so we split the ackIds
	Tuple2<List<String>, List<String>> splittedAckIds = splitAckIds(acknowledgementIds);
	while (!splittedAckIds.f0.isEmpty()) {
		AcknowledgeRequest acknowledgeRequest =
				AcknowledgeRequest.newBuilder()
								.setSubscription(projectSubscriptionName)
								.addAllAckIds(splittedAckIds.f0)
								.build();

		stub.withDeadlineAfter(60, SECONDS).acknowledge(acknowledgeRequest);

		splittedAckIds = splitAckIds(splittedAckIds.f1);
	}
}
 

@Override
public void acknowledge(List<String> acknowledgementIds) {
	if (acknowledgementIds.isEmpty()) {
		return;
	}

	//grpc servers won't accept acknowledge requests that are too large so we split the ackIds
	Tuple2<List<String>, List<String>> splittedAckIds = splitAckIds(acknowledgementIds);
	while (!splittedAckIds.f0.isEmpty()) {
		AcknowledgeRequest acknowledgeRequest =
				AcknowledgeRequest.newBuilder()
								.setSubscription(projectSubscriptionName)
								.addAllAckIds(splittedAckIds.f0)
								.build();

		stub.withDeadlineAfter(60, SECONDS).acknowledge(acknowledgeRequest);

		splittedAckIds = splitAckIds(splittedAckIds.f1);
	}
}
 

/**
 * Pull and ack {@link PubsubMessage}s from the subscription indicated by {@code index}.
 */
public List<PubsubMessage> pull(int index, int maxMessages, boolean returnImmediately) {
  List<ReceivedMessage> response = subscriber.pullCallable()
      .call(PullRequest.newBuilder().setMaxMessages(maxMessages)
          .setReturnImmediately(returnImmediately).setSubscription(getSubscription(index))
          .build())
      .getReceivedMessagesList();

  if (response.size() > 0) {
    subscriber.acknowledgeCallable()
        .call(AcknowledgeRequest.newBuilder().setSubscription(getSubscription(index))
            .addAllAckIds(
                response.stream().map(ReceivedMessage::getAckId).collect(Collectors.toList()))
            .build());
  }

  return response.stream().map(ReceivedMessage::getMessage).collect(Collectors.toList());
}
 

/** Tests that the no partition is assigned when the partition scheme is "kafka_partitioner". */
@Test
public void testPollWithPartitionSchemeKafkaPartitioner() throws Exception {
  props.put(
          CloudPubSubSourceConnector.KAFKA_PARTITION_SCHEME_CONFIG,
          CloudPubSubSourceConnector.PartitionScheme.KAFKA_PARTITIONER.toString());
  task.start(props);
  ReceivedMessage rm = createReceivedMessage(ACK_ID1, CPS_MESSAGE, new HashMap<String, String>());
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(1, result.size());
  SourceRecord expected =
          new SourceRecord(
                  null,
                  null,
                  KAFKA_TOPIC,
                  null,
                  Schema.OPTIONAL_STRING_SCHEMA,
                  null,
                  Schema.BYTES_SCHEMA,
                  KAFKA_VALUE);
  assertRecordsEqual(expected, result.get(0));
  assertNull(result.get(0).kafkaPartition());
}
 

/** Tests that the correct partition is assigned when the partition scheme is "hash_value". */
@Test
public void testPollWithPartitionSchemeHashValue() throws Exception {
  props.put(
      CloudPubSubSourceConnector.KAFKA_PARTITION_SCHEME_CONFIG,
      CloudPubSubSourceConnector.PartitionScheme.HASH_VALUE.toString());
  task.start(props);
  ReceivedMessage rm = createReceivedMessage(ACK_ID1, CPS_MESSAGE, new HashMap<String, String>());
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(1, result.size());
  SourceRecord expected =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          KAFKA_VALUE.hashCode() % Integer.parseInt(KAFKA_PARTITIONS),
          Schema.OPTIONAL_STRING_SCHEMA,
          null,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  assertRecordsEqual(expected, result.get(0));
}
 

/**
 * Tests when the message(s) retrieved from Cloud Pub/Sub do have an attribute that matches {@link
 * #KAFKA_MESSAGE_TIMESTAMP_ATTRIBUTE} and {@link #KAFKA_MESSAGE_KEY_ATTRIBUTE}.
 */
@Test
public void testPollWithMessageTimestampAttribute() throws Exception{
  task.start(props);
  Map<String, String> attributes = new HashMap<>();
  attributes.put(KAFKA_MESSAGE_KEY_ATTRIBUTE, KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE);
  attributes.put(KAFKA_MESSAGE_TIMESTAMP_ATTRIBUTE, KAFKA_MESSAGE_TIMESTAMP_ATTRIBUTE_VALUE);
  ReceivedMessage rm = createReceivedMessage(ACK_ID1, CPS_MESSAGE, attributes);
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(1, result.size());
  SourceRecord expected =
          new SourceRecord(
                  null,
                  null,
                  KAFKA_TOPIC,
                  0,
                  Schema.OPTIONAL_STRING_SCHEMA,
                  KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE,
                  Schema.BYTES_SCHEMA,
                  KAFKA_VALUE, Long.parseLong(KAFKA_MESSAGE_TIMESTAMP_ATTRIBUTE_VALUE));
  assertRecordsEqual(expected, result.get(0));
}
 

/**
 * Tests when the message(s) retrieved from Cloud Pub/Sub do have an attribute that matches {@link
 * #KAFKA_MESSAGE_KEY_ATTRIBUTE}.
 */
@Test
public void testPollWithMessageKeyAttribute() throws Exception {
  task.start(props);
  Map<String, String> attributes = new HashMap<>();
  attributes.put(KAFKA_MESSAGE_KEY_ATTRIBUTE, KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE);
  ReceivedMessage rm = createReceivedMessage(ACK_ID1, CPS_MESSAGE, attributes);
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(1, result.size());
  SourceRecord expected =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          0,
          Schema.OPTIONAL_STRING_SCHEMA,
          KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  assertRecordsEqual(expected, result.get(0));
}
 

/**
 * Tests when the message(s) retrieved from Cloud Pub/Sub do not have an attribute that matches
 * {@link #KAFKA_MESSAGE_KEY_ATTRIBUTE}.
 */
@Test
public void testPollWithNoMessageKeyAttribute() throws Exception {
  task.start(props);
  ReceivedMessage rm = createReceivedMessage(ACK_ID1, CPS_MESSAGE, new HashMap<String, String>());
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(1, result.size());
  SourceRecord expected =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          0,
          Schema.OPTIONAL_STRING_SCHEMA,
          null,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  assertRecordsEqual(expected, result.get(0));
}
 

/**
 * Tests that when ackMessages() succeeds and the subsequent call to poll() has no messages, that
 * the subscriber does not invoke ackMessages because there should be no acks.
 */
@Test
public void testPollInRegularCase() throws Exception {
  task.start(props);
  ReceivedMessage rm1 = createReceivedMessage(ACK_ID1, CPS_MESSAGE, new HashMap<String, String>());
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm1).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  assertEquals(1, result.size());
  task.commitRecord(result.get(0));
  stubbedPullResponse = PullResponse.newBuilder().build();
  SettableApiFuture<Empty> goodFuture = SettableApiFuture.create();
  goodFuture.set(Empty.getDefaultInstance());
  when(subscriber.ackMessages(any(AcknowledgeRequest.class))).thenReturn(goodFuture);
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  result = task.poll();
  assertEquals(0, result.size());
  result = task.poll();
  assertEquals(0, result.size());
  verify(subscriber, times(1)).ackMessages(any(AcknowledgeRequest.class));
}
 

public List<ReceivedMessage> pullMessages(String projectId, String subscriptionId, int maxNumberOfMessages) throws Exception {
	SubscriberStubSettings subscriberStubSettings =
		SubscriberStubSettings.newBuilder()
			.setTransportChannelProvider(channelProvider)
			.setCredentialsProvider(NoCredentialsProvider.create())
			.build();
	try (SubscriberStub subscriber = GrpcSubscriberStub.create(subscriberStubSettings)) {
		// String projectId = "my-project-id";
		// String subscriptionId = "my-subscription-id";
		// int numOfMessages = 10;   // max number of messages to be pulled
		String subscriptionName = ProjectSubscriptionName.format(projectId, subscriptionId);
		PullRequest pullRequest =
			PullRequest.newBuilder()
				.setMaxMessages(maxNumberOfMessages)
				.setReturnImmediately(false) // return immediately if messages are not available
				.setSubscription(subscriptionName)
				.build();

		// use pullCallable().futureCall to asynchronously perform this operation
		PullResponse pullResponse = subscriber.pullCallable().call(pullRequest);
		List<String> ackIds = new ArrayList<>();
		for (ReceivedMessage message : pullResponse.getReceivedMessagesList()) {
			// handle received message
			// ...
			ackIds.add(message.getAckId());
		}
		// acknowledge received messages
		AcknowledgeRequest acknowledgeRequest =
			AcknowledgeRequest.newBuilder()
				.setSubscription(subscriptionName)
				.addAllAckIds(ackIds)
				.build();
		// use acknowledgeCallable().futureCall to asynchronously perform this operation
		subscriber.acknowledgeCallable().call(acknowledgeRequest);
		return pullResponse.getReceivedMessagesList();
	}
}
 

private void acknowledgeIds(SubscriberStub subscriber, String subscriptionName, List<ReceivedMessage> receivedMessages) {
	if (receivedMessages.isEmpty()) {
		return;
	}

	List<String> ackIds = receivedMessages.stream().map(ReceivedMessage::getAckId).collect(Collectors.toList());
	// acknowledge received messages
	AcknowledgeRequest acknowledgeRequest =
		AcknowledgeRequest.newBuilder()
						.setSubscription(subscriptionName)
						.addAllAckIds(ackIds)
						.build();
	// use acknowledgeCallable().futureCall to asynchronously perform this operation
	subscriber.acknowledgeCallable().call(acknowledgeRequest);
}
 

@Override
public void acknowledge(SubscriptionPath subscription, List<String> ackIds) throws IOException {
  AcknowledgeRequest request =
      AcknowledgeRequest.newBuilder()
          .setSubscription(subscription.getPath())
          .addAllAckIds(ackIds)
          .build();
  subscriberStub().acknowledge(request); // ignore Empty result.
}
 

/**
 * Tests when the message retrieved from Cloud Pub/Sub have several attributes, including
 * one that matches {@link #KAFKA_MESSAGE_KEY_ATTRIBUTE}
 */
@Test
public void testPollWithMultipleAttributes() throws Exception {
  task.start(props);
  Map<String, String> attributes = new HashMap<>();
  attributes.put(KAFKA_MESSAGE_KEY_ATTRIBUTE, KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE);
  attributes.put("attribute1", "attribute_value1");
  attributes.put("attribute2", "attribute_value2");
  ReceivedMessage rm = createReceivedMessage(ACK_ID1, CPS_MESSAGE, attributes);
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(1, result.size());
  Schema expectedSchema =
      SchemaBuilder.struct()
          .field(ConnectorUtils.KAFKA_MESSAGE_CPS_BODY_FIELD, Schema.BYTES_SCHEMA)
          .field("attribute1", Schema.STRING_SCHEMA)
          .field("attribute2", Schema.STRING_SCHEMA)
          .build();
  Struct expectedValue = new Struct(expectedSchema)
                             .put(ConnectorUtils.KAFKA_MESSAGE_CPS_BODY_FIELD, KAFKA_VALUE)
                             .put("attribute1", "attribute_value1")
                             .put("attribute2", "attribute_value2");
  SourceRecord expected =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          0,
          Schema.OPTIONAL_STRING_SCHEMA,
          KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE,
          expectedSchema,
          expectedValue);
  assertRecordsEqual(expected, result.get(0));
}
 

/**
 * Tests when the message retrieved from Cloud Pub/Sub have several attributes, including
 * one that matches {@link #KAFKA_MESSAGE_KEY_ATTRIBUTE} and uses Kafka Record Headers to store them
 */
@Test
public void testPollWithMultipleAttributesAndRecordHeaders() throws Exception {
  props.put(CloudPubSubSourceConnector.USE_KAFKA_HEADERS, "true");
  task.start(props);
  Map<String, String> attributes = new HashMap<>();
  attributes.put(KAFKA_MESSAGE_KEY_ATTRIBUTE, KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE);
  attributes.put("attribute1", "attribute_value1");
  attributes.put("attribute2", "attribute_value2");
  ReceivedMessage rm = createReceivedMessage(ACK_ID1, CPS_MESSAGE, attributes);
  PullResponse stubbedPullResponse = PullResponse.newBuilder().addReceivedMessages(rm).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(1, result.size());

  ConnectHeaders headers = new ConnectHeaders();
  headers.addString("attribute1", "attribute_value1");
  headers.addString("attribute2", "attribute_value2");

  SourceRecord expected =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          0,
          Schema.OPTIONAL_STRING_SCHEMA,
          KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE,
          Long.parseLong(KAFKA_MESSAGE_TIMESTAMP_ATTRIBUTE_VALUE),
          headers);
  assertRecordsEqual(expected, result.get(0));
}
 

/** Tests when no messages are received from the Cloud Pub/Sub PullResponse. */
@Test
public void testPollCaseWithNoMessages() throws Exception {
  task.start(props);
  PullResponse stubbedPullResponse = PullResponse.newBuilder().build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  assertEquals(0, task.poll().size());
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
}
 

/**
 * Attempt to ack all ids in {@link #deliveredAckIds}.
 */
private void ackMessages() {
  if (deliveredAckIds.size() != 0) {
    AcknowledgeRequest.Builder requestBuilder = AcknowledgeRequest.newBuilder()
        .setSubscription(cpsSubscription);
    final Set<String> ackIdsBatch = new HashSet<>();
    synchronized (deliveredAckIds) {
      requestBuilder.addAllAckIds(deliveredAckIds);
      ackIdsInFlight.addAll(deliveredAckIds);
      ackIdsBatch.addAll(deliveredAckIds);
      deliveredAckIds.clear();
    }
    final ApiFuture<Empty> response = subscriber.ackMessages(requestBuilder.build());
    response.addListener(new Runnable() {
      @Override
      public void run() {
        try {
          response.get();
          log.trace("Successfully acked a set of messages. {}", ackIdsBatch.size());
        } catch (Exception e) {
          deliveredAckIds.addAll(ackIdsBatch);
          log.error("An exception occurred acking messages: " + e);
        } finally {
          ackIdsInFlight.removeAll(ackIdsBatch);
        }
      }
    }, ackExecutor);
  }
}
 

@Test
public void testPull_AndManualMultiSubscriptionAck()
		throws InterruptedException, ExecutionException, TimeoutException {
	ExecutorService mockExecutor = Mockito.mock(ExecutorService.class);
	this.pubSubSubscriberTemplate.setAckExecutor(mockExecutor);

	List<AcknowledgeablePubsubMessage> result1 = this.pubSubSubscriberTemplate.pull(
			"sub1", 1, true);
	List<AcknowledgeablePubsubMessage> result2 = this.pubSubSubscriberTemplate.pull(
			"sub2", 1, true);
	Set<AcknowledgeablePubsubMessage> combinedMessages = new HashSet<>(result1);
	combinedMessages.addAll(result2);

	assertThat(combinedMessages.size()).isEqualTo(2);

	TestListenableFutureCallback testListenableFutureCallback = new TestListenableFutureCallback();

	ListenableFuture<Void> listenableFuture = this.pubSubSubscriberTemplate.ack(combinedMessages);
	assertThat(listenableFuture).isNotNull();

	listenableFuture.addCallback(testListenableFutureCallback);
	listenableFuture.get(10L, TimeUnit.SECONDS);

	assertThat(listenableFuture.isDone()).isTrue();
	assertThat(testListenableFutureCallback.getThrowable()).isNull();
	verify(this.ackCallable, times(2)).futureCall(any(AcknowledgeRequest.class));
	verify(this.ackApiFuture, times(2)).addListener(any(), same(mockExecutor));
}
 

private ApiFuture<Empty> ack(String subscriptionName, Collection<String> ackIds) {
	AcknowledgeRequest acknowledgeRequest = AcknowledgeRequest.newBuilder()
			.addAllAckIds(ackIds)
			.setSubscription(subscriptionName)
			.build();
	return this.subscriberStub.acknowledgeCallable().futureCall(acknowledgeRequest);
}
 

@Test
public void acknowledge_subscriptionDoesNotExist() {
  expectedException.expect(StatusRuntimeException.class);
  expectedException.expectMessage(Status.NOT_FOUND.getCode().toString());

  AcknowledgeRequest request =
      AcknowledgeRequest.newBuilder()
          .setSubscription("projects/project-1/subscriptions/unknown-subscription")
          .addAllAckIds(Arrays.asList("0-0", "0-1", "0-2"))
          .build();
  blockingStub.acknowledge(request);
}
 

@Test
public void acknowledge() {
  List<String> ackIds = Arrays.asList("0-0", "0-1", "0-2");
  when(mockSubscriptionManager3.acknowledge(ackIds)).thenReturn(ackIds);

  AcknowledgeRequest request =
      AcknowledgeRequest.newBuilder()
          .setSubscription(TestHelpers.PROJECT2_SUBSCRIPTION3)
          .addAllAckIds(ackIds)
          .build();
  assertThat(blockingStub.acknowledge(request), Matchers.equalTo(Empty.getDefaultInstance()));

  verify(mockSubscriptionManager3).acknowledge(ackIds);
}
 

@Override
public void acknowledge(AcknowledgeRequest request, StreamObserver<Empty> responseObserver) {
  logger.atFine().log("Acknowledging messages %s", request);
  SubscriptionManager subscriptionManager = subscriptions.get(request.getSubscription());
  if (subscriptionManager == null) {
    String message = request.getSubscription() + " is not a valid Subscription";
    logger.atWarning().log(message);
    responseObserver.onError(Status.NOT_FOUND.withDescription(message).asException());
  } else {
    List<String> ackIds = subscriptionManager.acknowledge(request.getAckIdsList());
    logger.atFine().log("Successfully acknowledged %d messages", ackIds.size());
    responseObserver.onNext(Empty.getDefaultInstance());
    responseObserver.onCompleted();
  }
}
 

@Override
public ApiFuture<Empty> ackMessages(AcknowledgeRequest request) {
  currentSubscriberIndex = (currentSubscriberIndex + 1) % subscribers.size();
  return subscribers.get(currentSubscriberIndex).ackMessages(request);
}
 

public ApiFuture<Empty> ackMessages(AcknowledgeRequest request) {
  if (System.currentTimeMillis() > nextSubscriberResetTime) {
    makeSubscriber();
  }
  return subscriber.acknowledgeCallable().futureCall(request);
}
 

/**
 * Tests that the correct partition is assigned when the partition scheme is "hash_key". The test
 * has two cases, one where a key does exist and one where it does not.
 */
@Test
public void testPollWithPartitionSchemeHashKey() throws Exception {
  props.put(
      CloudPubSubSourceConnector.KAFKA_PARTITION_SCHEME_CONFIG,
      CloudPubSubSourceConnector.PartitionScheme.HASH_KEY.toString());
  task.start(props);
  Map<String, String> attributes = new HashMap<>();
  attributes.put(KAFKA_MESSAGE_KEY_ATTRIBUTE, KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE);
  ReceivedMessage withoutKey = createReceivedMessage(ACK_ID1, CPS_MESSAGE, new HashMap<String, String>());
  ReceivedMessage withKey = createReceivedMessage(ACK_ID2, CPS_MESSAGE, attributes);
  PullResponse stubbedPullResponse =
      PullResponse.newBuilder()
          .addReceivedMessages(0, withKey)
          .addReceivedMessages(1, withoutKey)
          .build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(2, result.size());
  SourceRecord expectedForMessageWithKey =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE.hashCode() % Integer.parseInt(KAFKA_PARTITIONS),
          Schema.OPTIONAL_STRING_SCHEMA,
          KAFKA_MESSAGE_KEY_ATTRIBUTE_VALUE,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  SourceRecord expectedForMessageWithoutKey =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          0,
          Schema.OPTIONAL_STRING_SCHEMA,
          null,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);

  assertRecordsEqual(expectedForMessageWithKey, result.get(0));
  assertArrayEquals((byte[])expectedForMessageWithoutKey.value(), (byte[])result.get(1).value());
}
 

/**
 * Tests that the correct partition is assigned when the partition scheme is "round_robin". The
 * tests makes sure to submit an approrpriate number of messages to poll() so that all partitions
 * in the round robin are hit once.
 */
@Test
public void testPollWithPartitionSchemeRoundRobin() throws Exception {
  task.start(props);
  ReceivedMessage rm1 = createReceivedMessage(ACK_ID1, CPS_MESSAGE, new HashMap<String, String>());
  ReceivedMessage rm2 = createReceivedMessage(ACK_ID2, CPS_MESSAGE, new HashMap<String, String>());
  ReceivedMessage rm3 = createReceivedMessage(ACK_ID3, CPS_MESSAGE, new HashMap<String, String>());
  ReceivedMessage rm4 = createReceivedMessage(ACK_ID4, CPS_MESSAGE, new HashMap<String, String>());
  PullResponse stubbedPullResponse =
      PullResponse.newBuilder()
          .addReceivedMessages(0, rm1)
          .addReceivedMessages(1, rm2)
          .addReceivedMessages(2, rm3)
          .addReceivedMessages(3, rm4)
          .build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  List<SourceRecord> result = task.poll();
  verify(subscriber, never()).ackMessages(any(AcknowledgeRequest.class));
  assertEquals(4, result.size());
  SourceRecord expected1 =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          0,
          Schema.OPTIONAL_STRING_SCHEMA,
          null,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  SourceRecord expected2 =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          1,
          Schema.OPTIONAL_STRING_SCHEMA,
          null,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  SourceRecord expected3 =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          2,
          Schema.OPTIONAL_STRING_SCHEMA,
          null,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  SourceRecord expected4 =
      new SourceRecord(
          null,
          null,
          KAFKA_TOPIC,
          0,
          Schema.OPTIONAL_STRING_SCHEMA,
          null,
          Schema.BYTES_SCHEMA,
          KAFKA_VALUE);
  assertRecordsEqual(expected1, result.get(0));
  assertRecordsEqual(expected2, result.get(1));
  assertRecordsEqual(expected3, result.get(2));
  assertRecordsEqual(expected4, result.get(3));
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    if (subscriber == null) {

        if (storedException.get() != null) {
            getLogger().error("Failed to create Google Cloud PubSub subscriber due to {}", new Object[]{storedException.get()});
        } else {
            getLogger().error("Google Cloud PubSub Subscriber was not properly created. Yielding the processor...");
        }

        context.yield();
        return;
    }

    final PullResponse pullResponse = subscriber.pullCallable().call(pullRequest);
    final List<String> ackIds = new ArrayList<>();

    for (ReceivedMessage message : pullResponse.getReceivedMessagesList()) {
        if (message.hasMessage()) {
            FlowFile flowFile = session.create();

            final Map<String, String> attributes = new HashMap<>();
            ackIds.add(message.getAckId());

            attributes.put(ACK_ID_ATTRIBUTE, message.getAckId());
            attributes.put(SERIALIZED_SIZE_ATTRIBUTE, String.valueOf(message.getSerializedSize()));
            attributes.put(MESSAGE_ID_ATTRIBUTE, message.getMessage().getMessageId());
            attributes.put(MSG_ATTRIBUTES_COUNT_ATTRIBUTE, String.valueOf(message.getMessage().getAttributesCount()));
            attributes.put(MSG_PUBLISH_TIME_ATTRIBUTE, String.valueOf(message.getMessage().getPublishTime().getSeconds()));
            attributes.putAll(message.getMessage().getAttributesMap());

            flowFile = session.putAllAttributes(flowFile, attributes);
            flowFile = session.write(flowFile, out -> out.write(message.getMessage().getData().toByteArray()));

            session.transfer(flowFile, REL_SUCCESS);
            session.getProvenanceReporter().receive(flowFile, getSubscriptionName(context));
        }
    }

    if (!ackIds.isEmpty()) {
        AcknowledgeRequest acknowledgeRequest = AcknowledgeRequest.newBuilder()
                .addAllAckIds(ackIds)
                .setSubscription(getSubscriptionName(context))
                .build();
        subscriber.acknowledgeCallable().call(acknowledgeRequest);
    }
}
 

public ApiFuture<Empty> ackMessages(AcknowledgeRequest request);