import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
/**
 * Tests that a call to flush() processes the Futures that were generated during a previous
 * call to put() (i.e enough messages were buffered in put() to trigger a publish).
 */
@Test
public void testFlushWithPublishInPut() throws Exception {
  props.put(CloudPubSubSinkConnector.MAX_BUFFER_SIZE_CONFIG, CPS_MIN_BATCH_SIZE1);
  task.start(props);
  List<SinkRecord> records = getSampleRecords();
  ListenableFuture<PublishResponse> goodFuture =
      spy(Futures.immediateFuture(PublishResponse.getDefaultInstance()));
  when(publisher.publish(any(PublishRequest.class))).thenReturn(goodFuture);
  task.put(records);
  Map<TopicPartition, OffsetAndMetadata> partitionOffsets = new HashMap<>();
  partitionOffsets.put(new TopicPartition(KAFKA_TOPIC, 0), null);
  task.flush(partitionOffsets);
  verify(goodFuture, times(1)).get();
  ArgumentCaptor<PublishRequest> captor = ArgumentCaptor.forClass(PublishRequest.class);
  verify(publisher, times(1)).publish(captor.capture());
  PublishRequest requestArg = captor.getValue();
  assertEquals(requestArg.getMessagesList(), getPubsubMessagesFromSampleRecords());
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
/** Publish all the messages in a partition and store the Future's for each publish request. */
private void publishMessagesForPartition(
    String topic, Integer partition, List<PubsubMessage> messages) {
  // Get a map containing all futures per partition for the passed in topic.
  Map<Integer, OutstandingFuturesForPartition> outstandingFuturesForTopic =
      allOutstandingFutures.get(topic);
  if (outstandingFuturesForTopic == null) {
    outstandingFuturesForTopic = new HashMap<>();
    allOutstandingFutures.put(topic, outstandingFuturesForTopic);
  }
  // Get the object containing the outstanding futures for this topic and partition..
  OutstandingFuturesForPartition outstandingFutures = outstandingFuturesForTopic.get(partition);
  if (outstandingFutures == null) {
    outstandingFutures = new OutstandingFuturesForPartition();
    outstandingFuturesForTopic.put(partition, outstandingFutures);
  }
  int startIndex = 0;
  int endIndex = Math.min(CPS_MAX_MESSAGES_PER_REQUEST, messages.size());
  PublishRequest.Builder builder = PublishRequest.newBuilder();
  // Publish all the messages for this partition in batches.
  while (startIndex < messages.size()) {
    PublishRequest request =
        builder.setTopic(cpsTopic).addAllMessages(messages.subList(startIndex, endIndex)).build();
    builder.clear();
    log.trace("Publishing: " + (endIndex - startIndex) + " messages");
    outstandingFutures.futures.add(publisher.publish(request));
    startIndex = endIndex;
    endIndex = Math.min(endIndex + CPS_MAX_MESSAGES_PER_REQUEST, messages.size());
  }
  messages.clear();
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
/**
 * Tests that if there are not enough messages buffered, publisher.publish() is not invoked.
 */
@Test
public void testPutWhereNoPublishesAreInvoked() {
  task.start(props);
  List<SinkRecord> records = getSampleRecords();
  task.put(records);
  verify(publisher, never()).publish(any(PublishRequest.class));
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
/**
 * Tests that if there are enough messages buffered, that the PublishRequest sent to the publisher
 * is correct.
 */
@Test
public void testPutWherePublishesAreInvoked() {
  props.put(CloudPubSubSinkConnector.MAX_BUFFER_SIZE_CONFIG, CPS_MIN_BATCH_SIZE1);
  task.start(props);
  List<SinkRecord> records = getSampleRecords();
  task.put(records);
  ArgumentCaptor<PublishRequest> captor = ArgumentCaptor.forClass(PublishRequest.class);
  verify(publisher, times(1)).publish(captor.capture());
  PublishRequest requestArg = captor.getValue();
  assertEquals(requestArg.getMessagesList(), getPubsubMessagesFromSampleRecords());
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
/**
 * Tests that a call to flush() processes the Futures that were generated during this same call
 * to flush() (i.e buffered messages were not published until the call to flush()).
 */
@Test
public void testFlushWithNoPublishInPut() throws Exception {
  task.start(props);
  Map<TopicPartition, OffsetAndMetadata> partitionOffsets = new HashMap<>();
  partitionOffsets.put(new TopicPartition(KAFKA_TOPIC, 0), null);
  List<SinkRecord> records = getSampleRecords();
  ListenableFuture<PublishResponse> goodFuture =
      spy(Futures.immediateFuture(PublishResponse.getDefaultInstance()));
  when(publisher.publish(any(PublishRequest.class))).thenReturn(goodFuture);
  task.put(records);
  task.flush(partitionOffsets);
  verify(publisher, times(1)).publish(any(PublishRequest.class));
  verify(goodFuture, times(1)).get();
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
/**
 * Tests that if a Future that is being processed in flush() failed with an exception, that an
 * exception is thrown.
 */
@Test(expected = RuntimeException.class)
public void testFlushExceptionCase() throws Exception {
  task.start(props);
  Map<TopicPartition, OffsetAndMetadata> partitionOffsets = new HashMap<>();
  partitionOffsets.put(new TopicPartition(KAFKA_TOPIC, 0), null);
  List<SinkRecord> records = getSampleRecords();
  ListenableFuture<PublishResponse> badFuture = spy(Futures.<PublishResponse>immediateFailedFuture(new Exception()));
  when(publisher.publish(any(PublishRequest.class))).thenReturn(badFuture);
  task.put(records);
  task.flush(partitionOffsets);
  verify(publisher, times(1)).publish(any(PublishRequest.class));
  verify(badFuture, times(1)).get();
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
@Override
public int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages)
    throws IOException {
  PublishRequest.Builder request = PublishRequest.newBuilder()
                                                 .setTopic(topic.getPath());
  for (OutgoingMessage outgoingMessage : outgoingMessages) {
    PubsubMessage.Builder message =
        PubsubMessage.newBuilder()
                     .setData(ByteString.copyFrom(outgoingMessage.elementBytes));

    if (outgoingMessage.attributes != null) {
      message.putAllAttributes(outgoingMessage.attributes);
    }

    if (timestampAttribute != null) {
      message.getMutableAttributes()
             .put(timestampAttribute, String.valueOf(outgoingMessage.timestampMsSinceEpoch));
    }

    if (idAttribute != null && !Strings.isNullOrEmpty(outgoingMessage.recordId)) {
      message.getMutableAttributes().put(idAttribute, outgoingMessage.recordId);
    }

    request.addMessages(message);
  }

  PublishResponse response = publisherStub().publish(request.build());
  return response.getMessageIdsCount();
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
@Override
public ListenableFuture<PublishResponse> publish(PublishRequest request) {
  return publisher.publish(request);
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
@Override
public ListenableFuture<PublishResponse> publish(PublishRequest request) {
  currentPublisherIndex = (currentPublisherIndex + 1) % publishers.size();
  return publishers.get(currentPublisherIndex).publish(request);
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
@Test
public void publishOneMessage() throws IOException {
  String expectedTopic = TOPIC.getPath();
  PubsubMessage expectedPubsubMessage =
      PubsubMessage.newBuilder()
                   .setData(ByteString.copyFrom(DATA.getBytes()))
                   .putAllAttributes(ATTRIBUTES)
                   .putAllAttributes(
                       ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME),
                                       ID_ATTRIBUTE, RECORD_ID))
                   .build();
  final PublishRequest expectedRequest =
      PublishRequest.newBuilder()
                    .setTopic(expectedTopic)
                    .addAllMessages(
                        ImmutableList.of(expectedPubsubMessage))
                    .build();
  final PublishResponse response =
      PublishResponse.newBuilder()
                     .addAllMessageIds(ImmutableList.of(MESSAGE_ID))
                     .build();

  final List<PublishRequest> requestsReceived = new ArrayList<>();
  PublisherImplBase publisherImplBase = new PublisherImplBase() {
    @Override
    public void publish(
        PublishRequest request, StreamObserver<PublishResponse> responseObserver) {
      requestsReceived.add(request);
      responseObserver.onNext(response);
      responseObserver.onCompleted();
    }
  };
  Server server = InProcessServerBuilder.forName(channelName)
      .addService(publisherImplBase)
      .build()
      .start();
  try {
    OutgoingMessage actualMessage = new OutgoingMessage(
            DATA.getBytes(), ATTRIBUTES, MESSAGE_TIME, RECORD_ID);
    int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
    assertEquals(1, n);
    assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
  } finally {
    server.shutdownNow();
  }
}
 

import com.google.pubsub.v1.PublishRequest; //导入依赖的package包/类
public ListenableFuture<PublishResponse> publish(PublishRequest request);