@Override
  public List<PartitionGroupMetadata> computePartitionGroupMetadata(String clientId, StreamConfig streamConfig,
      List<PartitionGroupConsumptionStatus> partitionGroupConsumptionStatuses, int timeoutMillis)
      throws IOException, TimeoutException {

    List<PartitionGroupMetadata> newPartitionGroupMetadataList = new ArrayList<>();

    Map<String, Shard> shardIdToShardMap = _kinesisConnectionHandler.getShards().stream()
        .collect(Collectors.toMap(Shard::shardId, s -> s, (s1, s2) -> s1));
    Set<String> shardsInCurrent = new HashSet<>();
    Set<String> shardsEnded = new HashSet<>();

    // TODO: Once we start supporting multiple shards in a PartitionGroup,
    //  we need to iterate over all shards to check if any of them have reached end

    // Process existing shards. Add them to new list if still consuming from them
    for (PartitionGroupConsumptionStatus currentPartitionGroupConsumptionStatus : partitionGroupConsumptionStatuses) {
      KinesisPartitionGroupOffset kinesisStartCheckpoint =
          (KinesisPartitionGroupOffset) currentPartitionGroupConsumptionStatus.getStartOffset();
      String shardId = kinesisStartCheckpoint.getShardId();
      shardsInCurrent.add(shardId);
      Shard shard = shardIdToShardMap.get(shardId);
      if (shard == null) { // Shard has expired
        shardsEnded.add(shardId);
        String lastConsumedSequenceID = kinesisStartCheckpoint.getSequenceNumber();
        LOGGER.warn(
            "Kinesis shard with id: {} has expired. Data has been consumed from the shard till sequence number: {}. "
                + "There can be potential data loss.", shardId, lastConsumedSequenceID);
        continue;
      }

      StreamPartitionMsgOffset newStartOffset;
      StreamPartitionMsgOffset currentEndOffset = currentPartitionGroupConsumptionStatus.getEndOffset();
      if (currentEndOffset != null) { // Segment DONE (committing/committed)
        String endingSequenceNumber = shard.sequenceNumberRange().endingSequenceNumber();
        if (endingSequenceNumber != null) { // Shard has ended, check if we're also done consuming it
          if (consumedEndOfShard(currentEndOffset, currentPartitionGroupConsumptionStatus)) {
            shardsEnded.add(shardId);
            continue; // Shard ended and we're done consuming it. Skip
          }
        }
        newStartOffset = currentEndOffset;
      } else { // Segment IN_PROGRESS
        newStartOffset = currentPartitionGroupConsumptionStatus.getStartOffset();
      }
      newPartitionGroupMetadataList.add(
          new PartitionGroupMetadata(currentPartitionGroupConsumptionStatus.getPartitionGroupId(), newStartOffset));
    }

    // Add brand new shards
    for (Map.Entry<String, Shard> entry : shardIdToShardMap.entrySet()) {
      // If shard was already in current list, skip
      String newShardId = entry.getKey();
      if (shardsInCurrent.contains(newShardId)) {
        continue;
      }
      Shard newShard = entry.getValue();
      String parentShardId = newShard.parentShardId();

      // Add the new shard in the following 3 cases:
      // 1. Root shards - Parent shardId will be null. Will find this case when creating new table.
      // 2. Parent expired - Parent shardId will not be part of shardIdToShard map
      // 3. Parent reached EOL and completely consumed.
      if (parentShardId == null || !shardIdToShardMap.containsKey(parentShardId) || shardsEnded.contains(
          parentShardId)) {
        // TODO: Revisit this. Kinesis starts consuming AFTER the start sequence number, and we might miss the first
        //       message.
        StreamPartitionMsgOffset newStartOffset =
            new KinesisPartitionGroupOffset(newShardId, newShard.sequenceNumberRange().startingSequenceNumber());
        int partitionGroupId = getPartitionGroupIdFromShardId(newShardId);
        newPartitionGroupMetadataList.add(new PartitionGroupMetadata(partitionGroupId, newStartOffset));
      }
    }
    return newPartitionGroupMetadataList;
  }
@Test
  public void getPartitionsGroupInfoChildShardsest()
      throws Exception {
    List<PartitionGroupConsumptionStatus> currentPartitionGroupMeta = new ArrayList<>();

    Map<String, String> shardToSequenceMap = new HashMap<>();
    shardToSequenceMap.put("1", "1");
    KinesisPartitionGroupOffset kinesisPartitionGroupOffset = new KinesisPartitionGroupOffset("1", "1");

    currentPartitionGroupMeta.add(
        new PartitionGroupConsumptionStatus(0, 1, kinesisPartitionGroupOffset, kinesisPartitionGroupOffset,
            "CONSUMING"));

    ArgumentCaptor<StreamPartitionMsgOffset> checkpointArgs = ArgumentCaptor.forClass(StreamPartitionMsgOffset.class);
    ArgumentCaptor<PartitionGroupConsumptionStatus> partitionGroupMetadataCapture =
        ArgumentCaptor.forClass(PartitionGroupConsumptionStatus.class);
    ArgumentCaptor<Integer> intArguments = ArgumentCaptor.forClass(Integer.class);
    ArgumentCaptor<String> stringCapture = ArgumentCaptor.forClass(String.class);

    Shard shard0 = Shard.builder().shardId(SHARD_ID_0).parentShardId(SHARD_ID_1)
        .sequenceNumberRange(SequenceNumberRange.builder().startingSequenceNumber("1").build()).build();
    Shard shard1 = Shard.builder().shardId(SHARD_ID_1).sequenceNumberRange(
        SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("1").build()).build();

    when(_kinesisConnectionHandler.getShards()).thenReturn(ImmutableList.of(shard0, shard1));
    when(_streamConsumerFactory.createPartitionGroupConsumer(stringCapture.capture(),
        partitionGroupMetadataCapture.capture())).thenReturn(_partitionGroupConsumer);
    when(_partitionGroupConsumer.fetchMessages(checkpointArgs.capture(), intArguments.capture())).thenReturn(
        new KinesisMessageBatch(new ArrayList<>(), kinesisPartitionGroupOffset, true));

    List<PartitionGroupMetadata> result =
        _kinesisStreamMetadataProvider.computePartitionGroupMetadata(CLIENT_ID, getStreamConfig(),
            currentPartitionGroupMeta, TIMEOUT);

    Assert.assertEquals(result.size(), 1);
    Assert.assertEquals(result.get(0).getPartitionGroupId(), 0);
    Assert.assertEquals(partitionGroupMetadataCapture.getValue().getSequenceNumber(), 1);
  }