com.google.pubsub.v1.ReceivedMessage Java Examples

void processMessage(SourceContext<OUT> sourceContext, List<ReceivedMessage> messages) throws Exception {
	synchronized (sourceContext.getCheckpointLock()) {
		for (ReceivedMessage message : messages) {
			acknowledgeOnCheckpoint.addAcknowledgeId(message.getAckId());

			PubsubMessage pubsubMessage = message.getMessage();

			OUT deserializedMessage = deserializationSchema.deserialize(pubsubMessage);
			if (deserializationSchema.isEndOfStream(deserializedMessage)) {
				cancel();
				return;
			}

			sourceContext.collect(deserializedMessage);
		}

	}
}
 

/** 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 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 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 a call to ackMessages() fails, that the message is not redelivered to Kafka if
 * the message is received again by Cloud Pub/Sub. Also tests that ack ids are added properly if
 * the ack id has not been seen before.
 */
@Test
public void testPollWithDuplicateReceivedMessages() 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());
  ReceivedMessage rm2 = createReceivedMessage(ACK_ID2, CPS_MESSAGE, new HashMap<String, String>());
  stubbedPullResponse =
      PullResponse.newBuilder().addReceivedMessages(0, rm1).addReceivedMessages(1, rm2).build();
  when(subscriber.pull(any(PullRequest.class)).get()).thenReturn(stubbedPullResponse);
  result = task.poll();
  assertEquals(1, result.size());
}
 

/**
 * 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));
}
 

/**
 * 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());
}
 

private void processMessage(
		SourceContext<OUT> sourceContext,
		List<ReceivedMessage> messages,
		PubSubCollector collector) throws Exception {
	rateLimiter.acquire(messages.size());

	synchronized (sourceContext.getCheckpointLock()) {
		for (ReceivedMessage message : messages) {
			acknowledgeOnCheckpoint.addAcknowledgeId(message.getAckId());

			PubsubMessage pubsubMessage = message.getMessage();

			deserializationSchema.deserialize(pubsubMessage, collector);
			if (collector.isEndOfStreamSignalled()) {
				cancel();
				return;
			}
		}

	}
}
 

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 subscriptionName = ProjectSubscriptionName.format(projectId, subscriptionId);
		PullRequest pullRequest =
			PullRequest.newBuilder()
				.setMaxMessages(maxNumberOfMessages)
				.setReturnImmediately(false)
				.setSubscription(subscriptionName)
				.build();

		List<ReceivedMessage> receivedMessages = subscriber.pullCallable().call(pullRequest).getReceivedMessagesList();
		acknowledgeIds(subscriber, subscriptionName, receivedMessages);
		return receivedMessages;
	}
}
 

@Test
public void pull() {
  List<PubsubMessage> messages =
      Arrays.asList(
          PubsubMessage.newBuilder()
              .setMessageId("0-0")
              .setData(ByteString.copyFromUtf8("hello"))
              .build(),
          PubsubMessage.newBuilder()
              .setMessageId("0-1")
              .setData(ByteString.copyFromUtf8("world"))
              .build());
  when(mockSubscriptionManager3.pull(100, false)).thenReturn(messages);

  PullRequest request =
      PullRequest.newBuilder()
          .setSubscription(TestHelpers.PROJECT2_SUBSCRIPTION3)
          .setMaxMessages(100)
          .build();
  PullResponse response = blockingStub.pull(request);
  assertThat(
      response.getReceivedMessagesList(),
      Matchers.contains(
          ReceivedMessage.newBuilder().setAckId("0-0").setMessage(messages.get(0)).build(),
          ReceivedMessage.newBuilder().setAckId("0-1").setMessage(messages.get(1)).build()));
}
 

@Test
public void testStoppingConnectorWhenDeserializationSchemaIndicatesEndOfStream() throws Exception {
	when(deserializationSchema.deserialize(pubSubMessage(FIRST_MESSAGE))).thenReturn(FIRST_MESSAGE);
	when(sourceContext.getCheckpointLock()).thenReturn("some object to lock on");
	pubSubSource.open(null);

	when(deserializationSchema.isEndOfStream(FIRST_MESSAGE)).thenReturn(true);

	ReceivedMessage message = receivedMessage("ackId", pubSubMessage(FIRST_MESSAGE));

	//Process message
	pubSubSource.processMessage(sourceContext, singletonList(message));
	verify(deserializationSchema, times(1)).isEndOfStream(FIRST_MESSAGE);
	verify(sourceContext, times(1)).getCheckpointLock();
	verifyNoMoreInteractions(sourceContext);
}
 

@Test
public void testPull() throws Exception {
	Publisher publisher = pubsubHelper.createPublisher(PROJECT_NAME, TOPIC_NAME);
	publisher
		.publish(PubsubMessage
			.newBuilder()
			.setData(ByteString.copyFromUtf8("Hello World PULL"))
			.build())
		.get();

	List<ReceivedMessage> receivedMessages = pubsubHelper.pullMessages(PROJECT_NAME, SUBSCRIPTION_NAME, 1);

	assertEquals(1, receivedMessages.size());
	assertEquals("Hello World PULL", receivedMessages.get(0).getMessage().getData().toStringUtf8());

	publisher.shutdown();
}
 

@Test
public void testPull() throws Exception {
	Publisher publisher = pubsubHelper.createPublisher(PROJECT_NAME, TOPIC_NAME);
	publisher
		.publish(PubsubMessage
			.newBuilder()
			.setData(ByteString.copyFromUtf8("Hello World PULL"))
			.build())
		.get();

	List<ReceivedMessage> receivedMessages = pubsubHelper.pullMessages(PROJECT_NAME, SUBSCRIPTION_NAME, 10);
	assertEquals(1, receivedMessages.size());
	assertEquals("Hello World PULL", receivedMessages.get(0).getMessage().getData().toStringUtf8());

	publisher.shutdown();
}
 

private List<ReceivedMessage> buildReceivedMessageList(
    String subscriptionName, List<PubsubMessage> pubsubMessages) {
  return pubsubMessages
      .stream()
      .map(
          m -> {
            statisticsManager.computeSubscriber(
                subscriptionName, m.getData(), m.getPublishTime());
            return ReceivedMessage.newBuilder().setAckId(m.getMessageId()).setMessage(m).build();
          })
      .collect(Collectors.toList());
}
 

private List<ReceivedMessage> pull(int retriesRemaining) {
	try {
		return stub.withDeadlineAfter(timeout.toMillis(), TimeUnit.MILLISECONDS)
				.pull(pullRequest)
				.getReceivedMessagesList();
	} catch (StatusRuntimeException e) {
		if (retriesRemaining > 0) {
			return pull(retriesRemaining - 1);
		}

		throw e;
	}
}
 

/**
 * 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));
}
 

@Test
public void testFlinkSink() throws Exception {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

	List<String> input = Arrays.asList("One", "Two", "Three", "Four", "Five", "Six", "Seven", "Eigth", "Nine", "Ten");

	// Create test stream
	DataStream<String> theData = env
		.fromCollection(input)
		.name("Test input")
		.map((MapFunction<String, String>) StringUtils::reverse);

	// Sink into pubsub
	theData
		.addSink(PubSubSink.newBuilder()
							.withSerializationSchema(new SimpleStringSchema())
							.withProjectName(PROJECT_NAME)
							.withTopicName(TOPIC_NAME)
						   // Specific for emulator
						.withHostAndPortForEmulator(getPubSubHostPort())
						.withCredentials(NoCredentials.getInstance())
						.build())
		.name("PubSub sink");

	// Run
	env.execute();

	// Now get the result from PubSub and verify if everything is there
	List<ReceivedMessage> receivedMessages = pubsubHelper.pullMessages(PROJECT_NAME, SUBSCRIPTION_NAME, 100);

	assertEquals("Wrong number of elements", input.size(), receivedMessages.size());

	// Check output strings
	List<String> output = new ArrayList<>();
	receivedMessages.forEach(msg -> output.add(msg.getMessage().getData().toStringUtf8()));

	for (String test : input) {
		assertTrue("Missing " + test, output.contains(StringUtils.reverse(test)));
	}
}
 

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 List<AcknowledgeablePubsubMessage> toAcknowledgeablePubsubMessageList(List<ReceivedMessage> messages,
		String subscriptionId) {
	return messages.stream()
			.map((message) -> new PulledAcknowledgeablePubsubMessage(
					PubSubSubscriptionUtils.toProjectSubscriptionName(subscriptionId,
							this.subscriberFactory.getProjectId()),
					message.getMessage(),
					message.getAckId()))
			.collect(Collectors.toList());
}
 

private List<ReceivedMessage> pull(int retriesRemaining) {
	try {
		return stub.withDeadlineAfter(timeout.toMillis(), TimeUnit.MILLISECONDS)
				.pull(pullRequest)
				.getReceivedMessagesList();
	} catch (StatusRuntimeException e) {
		if (retriesRemaining > 0) {
			return pull(retriesRemaining - 1);
		}

		throw e;
	}
}
 

@Test
public void testFlinkSink() throws Exception {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

	List<String> input = Arrays.asList("One", "Two", "Three", "Four", "Five", "Six", "Seven", "Eigth", "Nine", "Ten");

	// Create test stream
	DataStream<String> theData = env
		.fromCollection(input)
		.name("Test input")
		.map((MapFunction<String, String>) StringUtils::reverse);

	// Sink into pubsub
	theData
		.addSink(PubSubSink.newBuilder()
							.withSerializationSchema(new SimpleStringSchema())
							.withProjectName(PROJECT_NAME)
							.withTopicName(TOPIC_NAME)
						   // Specific for emulator
						.withHostAndPortForEmulator(getPubSubHostPort())
						.withCredentials(NoCredentials.getInstance())
						.build())
		.name("PubSub sink");

	// Run
	env.execute();

	// Now get the result from PubSub and verify if everything is there
	List<ReceivedMessage> receivedMessages = pubsubHelper.pullMessages(PROJECT_NAME, SUBSCRIPTION_NAME, 100);

	assertEquals("Wrong number of elements", input.size(), receivedMessages.size());

	// Check output strings
	List<String> output = new ArrayList<>();
	receivedMessages.forEach(msg -> output.add(msg.getMessage().getData().toStringUtf8()));

	for (String test : input) {
		assertTrue("Missing " + test, output.contains(StringUtils.reverse(test)));
	}
}
 

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 List<ReceivedMessage> pull() {
	ReceivedMessage message = queue.poll();
	if (message != null) {
		return Collections.singletonList(message);
	} else {
		return Collections.emptyList();
	}
}
 

@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);
    }
}
 

@Test
public void pullOneMessage() throws IOException {
  String expectedSubscription = SUBSCRIPTION.getPath();
  final PullRequest expectedRequest =
      PullRequest.newBuilder()
          .setSubscription(expectedSubscription)
          .setReturnImmediately(true)
          .setMaxMessages(10)
          .build();
  Timestamp timestamp =
      Timestamp.newBuilder()
          .setSeconds(PUB_TIME / 1000)
          .setNanos((int) (PUB_TIME % 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), 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> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
    assertEquals(1, acutalMessages.size());
    IncomingMessage actualMessage = acutalMessages.get(0);
    assertEquals(ACK_ID, actualMessage.ackId());
    assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
    assertEquals(RECORD_ID, actualMessage.recordId());
    assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch());
    assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch());
    assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
  } finally {
    server.shutdownNow();
  }
}
 

@Override
public List<ReceivedMessage> pull() {
	return pull(retries);
}
 

private ReceivedMessage receivedMessage(String ackId, PubsubMessage pubsubMessage) {
	return ReceivedMessage.newBuilder()
		.setAckId(ackId)
		.setMessage(pubsubMessage)
		.build();
}