// Path: src/test/java/com/google/cloud/pubsublite/spark/TestingUtils.java
// public static PslSourceOffset createPslSourceOffset(long... offsets) {
//   Map<Partition, Offset> map = new HashMap<>();
//   int idx = 0;
//   for (long offset : offsets) {
//     map.put(Partition.of(idx++), Offset.of(offset));
//   }
//   return PslSourceOffset.builder().partitionOffsetMap(map).build();
// }
// 
// Path: src/test/java/com/google/cloud/pubsublite/spark/TestingUtils.java
// public static SparkSourceOffset createSparkSourceOffset(long... offsets) {
//   Map<Partition, SparkPartitionOffset> map = new HashMap<>();
//   int idx = 0;
//   for (long offset : offsets) {
//     map.put(Partition.of(idx), SparkPartitionOffset.create(Partition.of(idx), offset));
//     idx++;
//   }
//   return new SparkSourceOffset(map);
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/MultiPartitionCommitter.java
// public interface MultiPartitionCommitter extends Closeable {
// 
//   interface CommitterFactory {
//     Committer newCommitter(Partition partition);
//   }
// 
//   void commit(PslSourceOffset offset);
// 
//   void close();
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/PartitionSubscriberFactory.java
// public interface PartitionSubscriberFactory extends Serializable {
//   Subscriber newSubscriber(
//       Partition partition, Offset offset, Consumer<List<SequencedMessage>> message_consumer)
//       throws ApiException;
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/PerTopicHeadOffsetReader.java
// public interface PerTopicHeadOffsetReader extends Closeable {
// 
//   // Gets the head offsets for all partitions in the topic. Blocks.
//   PslSourceOffset getHeadOffset();
// 
//   @Override
//   void close();
// }


import com.google.cloud.pubsublite.spark.internal.PerTopicHeadOffsetReader;
import com.google.common.collect.ImmutableMap;
import java.util.Optional;
import org.junit.Test;
import static com.google.cloud.pubsublite.spark.TestingUtils.createPslSourceOffset;
import static com.google.cloud.pubsublite.spark.TestingUtils.createSparkSourceOffset;
import static com.google.common.truth.Truth.assertThat;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.*;
import static org.mockito.Mockito.times;
import com.google.api.core.ApiFutures;
import com.google.cloud.pubsublite.Offset;
import com.google.cloud.pubsublite.Partition;
import com.google.cloud.pubsublite.internal.CursorClient;
import com.google.cloud.pubsublite.internal.testing.UnitTestExamples;
import com.google.cloud.pubsublite.spark.internal.MultiPartitionCommitter;
import com.google.cloud.pubsublite.spark.internal.PartitionSubscriberFactory;
    assertThat(reader.planInputPartitions()).hasSize(3);
  }

  @Test
  public void testPartitionIncreasedNewQuery() {
    when(cursorClient.listPartitionCursors(UnitTestExamples.exampleSubscriptionPath()))
        .thenReturn(ApiFutures.immediateFuture(ImmutableMap.of(Partition.of(0L), Offset.of(100L))));
    SparkSourceOffset endOffset = createSparkSourceOffset(301L, 200L);
    when(headOffsetReader.getHeadOffset()).thenReturn(PslSparkUtils.toPslSourceOffset(endOffset));
    reader.setOffsetRange(Optional.empty(), Optional.empty());
    assertThat(reader.getStartOffset()).isEqualTo(createSparkSourceOffset(99L, -1L));
    assertThat(reader.getEndOffset()).isEqualTo(endOffset);
    assertThat(reader.planInputPartitions()).hasSize(2);
  }

  @Test
  public void testPartitionIncreasedBeforeSetOffsets() {
    SparkSourceOffset endOffset = createSparkSourceOffset(301L, 200L);
    SparkSourceOffset startOffset = createSparkSourceOffset(100L);
    when(headOffsetReader.getHeadOffset()).thenReturn(PslSparkUtils.toPslSourceOffset(endOffset));
    reader.setOffsetRange(Optional.of(startOffset), Optional.empty());
    assertThat(reader.getStartOffset()).isEqualTo(startOffset);
    assertThat(reader.getEndOffset()).isEqualTo(endOffset);
    assertThat(reader.planInputPartitions()).hasSize(2);
  }

  @Test
  public void testPartitionIncreasedBetweenSetOffsetsAndPlan() {
    SparkSourceOffset startOffset = createSparkSourceOffset(100L);
    SparkSourceOffset endOffset = createSparkSourceOffset(301L);
// Path: src/test/java/com/google/cloud/pubsublite/spark/TestingUtils.java
// public static PslSourceOffset createPslSourceOffset(long... offsets) {
//   Map<Partition, Offset> map = new HashMap<>();
//   int idx = 0;
//   for (long offset : offsets) {
//     map.put(Partition.of(idx++), Offset.of(offset));
//   }
//   return PslSourceOffset.builder().partitionOffsetMap(map).build();
// }
// 
// Path: src/test/java/com/google/cloud/pubsublite/spark/TestingUtils.java
// public static SparkSourceOffset createSparkSourceOffset(long... offsets) {
//   Map<Partition, SparkPartitionOffset> map = new HashMap<>();
//   int idx = 0;
//   for (long offset : offsets) {
//     map.put(Partition.of(idx), SparkPartitionOffset.create(Partition.of(idx), offset));
//     idx++;
//   }
//   return new SparkSourceOffset(map);
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/MultiPartitionCommitter.java
// public interface MultiPartitionCommitter extends Closeable {
// 
//   interface CommitterFactory {
//     Committer newCommitter(Partition partition);
//   }
// 
//   void commit(PslSourceOffset offset);
// 
//   void close();
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/PartitionSubscriberFactory.java
// public interface PartitionSubscriberFactory extends Serializable {
//   Subscriber newSubscriber(
//       Partition partition, Offset offset, Consumer<List<SequencedMessage>> message_consumer)
//       throws ApiException;
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/PerTopicHeadOffsetReader.java
// public interface PerTopicHeadOffsetReader extends Closeable {
// 
//   // Gets the head offsets for all partitions in the topic. Blocks.
//   PslSourceOffset getHeadOffset();
// 
//   @Override
//   void close();
// }


// Path: src/test/java/com/google/cloud/pubsublite/spark/PslMicroBatchReaderTest.java
import com.google.cloud.pubsublite.spark.internal.PerTopicHeadOffsetReader;
import com.google.common.collect.ImmutableMap;
import java.util.Optional;
import org.junit.Test;
import static com.google.cloud.pubsublite.spark.TestingUtils.createPslSourceOffset;
import static com.google.cloud.pubsublite.spark.TestingUtils.createSparkSourceOffset;
import static com.google.common.truth.Truth.assertThat;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.*;
import static org.mockito.Mockito.times;
import com.google.api.core.ApiFutures;
import com.google.cloud.pubsublite.Offset;
import com.google.cloud.pubsublite.Partition;
import com.google.cloud.pubsublite.internal.CursorClient;
import com.google.cloud.pubsublite.internal.testing.UnitTestExamples;
import com.google.cloud.pubsublite.spark.internal.MultiPartitionCommitter;
import com.google.cloud.pubsublite.spark.internal.PartitionSubscriberFactory;

    assertThat(reader.planInputPartitions()).hasSize(3);
  }

  @Test
  public void testPartitionIncreasedNewQuery() {
    when(cursorClient.listPartitionCursors(UnitTestExamples.exampleSubscriptionPath()))
        .thenReturn(ApiFutures.immediateFuture(ImmutableMap.of(Partition.of(0L), Offset.of(100L))));
    SparkSourceOffset endOffset = createSparkSourceOffset(301L, 200L);
    when(headOffsetReader.getHeadOffset()).thenReturn(PslSparkUtils.toPslSourceOffset(endOffset));
    reader.setOffsetRange(Optional.empty(), Optional.empty());
    assertThat(reader.getStartOffset()).isEqualTo(createSparkSourceOffset(99L, -1L));
    assertThat(reader.getEndOffset()).isEqualTo(endOffset);
    assertThat(reader.planInputPartitions()).hasSize(2);
  }

  @Test
  public void testPartitionIncreasedBeforeSetOffsets() {
    SparkSourceOffset endOffset = createSparkSourceOffset(301L, 200L);
    SparkSourceOffset startOffset = createSparkSourceOffset(100L);
    when(headOffsetReader.getHeadOffset()).thenReturn(PslSparkUtils.toPslSourceOffset(endOffset));
    reader.setOffsetRange(Optional.of(startOffset), Optional.empty());
    assertThat(reader.getStartOffset()).isEqualTo(startOffset);
    assertThat(reader.getEndOffset()).isEqualTo(endOffset);
    assertThat(reader.planInputPartitions()).hasSize(2);
  }

  @Test
  public void testPartitionIncreasedBetweenSetOffsetsAndPlan() {
    SparkSourceOffset startOffset = createSparkSourceOffset(100L);
    SparkSourceOffset endOffset = createSparkSourceOffset(301L);
    SparkSourceOffset newEndOffset = createSparkSourceOffset(600L, 300L);