software.amazon.awssdk.services.kinesis.model.Shard Java Examples

/**
 * Helper method to create a new Lease POJO for a shard.
 * Note: Package level access only for testing purposes
 * 
 * @param shard
 * @return
 */
private static Lease newKCLLease(final Shard shard) {
    Lease newLease = new Lease();
    newLease.leaseKey(shard.shardId());
    List<String> parentShardIds = new ArrayList<>(2);
    if (shard.parentShardId() != null) {
        parentShardIds.add(shard.parentShardId());
    }
    if (shard.adjacentParentShardId() != null) {
        parentShardIds.add(shard.adjacentParentShardId());
    }
    newLease.parentShardIds(parentShardIds);
    newLease.ownerSwitchesSinceCheckpoint(0L);

    return newLease;
}
 

/**
 * Test determineNewLeasesToCreate() where there are no leases and no resharding operations have been performed, but
 * one of the shards was marked as inconsistent.
 */
@Test
public void testDetermineNewLeasesToCreate0Leases0Reshards1Inconsistent() {
    final String shardId0 = "shardId-0";
    final String shardId1 = "shardId-1";
    final String shardId2 = "shardId-2";
    final SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null);

    final List<Shard> shards = Arrays.asList(ShardObjectHelper.newShard(shardId0, null, null, sequenceRange),
            ShardObjectHelper.newShard(shardId1, null, null, sequenceRange),
            ShardObjectHelper.newShard(shardId2, shardId1, null, sequenceRange));
    final List<Lease> currentLeases = Collections.emptyList();

    final Set<String> inconsistentShardIds = new HashSet<>(Collections.singletonList(shardId2));

    final List<Lease> newLeases = HierarchicalShardSyncer.determineNewLeasesToCreate(shards, currentLeases,
            INITIAL_POSITION_LATEST, inconsistentShardIds);
    final Set<String> newLeaseKeys = newLeases.stream().map(Lease::leaseKey).collect(Collectors.toSet());
    final Set<String> expectedLeaseShardIds = new HashSet<>(Arrays.asList(shardId0, shardId1));
    assertThat(newLeases.size(), equalTo(expectedLeaseShardIds.size()));
    assertThat(newLeaseKeys, equalTo(expectedLeaseShardIds));
}
 

private List<Shard> mergeShards(final String streamName,
                                final String shard1,
                                final String shard2)
        throws InterruptedException {

    client.mergeShards(MergeShardsRequest.builder()
                                         .streamName(streamName)
                                         .shardToMerge(shard1)
                                         .adjacentShardToMerge(shard2)
                                         .build());

    List<Shard> shards = waitForStream(streamName);

    Assert.assertEquals(4, shards.size());
    Shard merged = shards.get(3);

    BigInteger start =
            new BigInteger(merged.hashKeyRange().startingHashKey());
    BigInteger end =
            new BigInteger(merged.hashKeyRange().endingHashKey());

    Assert.assertEquals(BigInteger.valueOf(0), start);
    Assert.assertTrue(end.compareTo(BigInteger.valueOf(1000)) >= 0);

    return shards;
}
 

@Test(expected = KinesisClientLibIOException.class)
public void testCheckAndCreateLeasesForNewShardsWhenParentIsOpen() throws Exception {
    final List<Shard> shards = new ArrayList<>(constructShardListForGraphA());
    final SequenceNumberRange range = shards.get(0).sequenceNumberRange().toBuilder().endingSequenceNumber(null)
            .build();
    final Shard shard = shards.get(3).toBuilder().sequenceNumberRange(range).build();
    shards.remove(3);
    shards.add(3, shard);

    when(shardDetector.listShards()).thenReturn(shards);

    try {
        hierarchicalShardSyncer.checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher,
                INITIAL_POSITION_TRIM_HORIZON, cleanupLeasesOfCompletedShards, false, SCOPE);
    } finally {
        verify(shardDetector).listShards();
        verify(dynamoDBLeaseRefresher, never()).listLeases();
    }
}
 

/**
 * Test CheckIfDescendantAndAddNewLeasesForAncestors - two parents, two ancestors, not descendant
 */
@Test
public void testCheckIfDescendantAndAddNewLeasesForAncestors2P2ANotDescendant() {
    final String parentShardId = "shardId-parent";
    final String adjacentParentShardId = "shardId-adjacentParent";
    final String shardId = "shardId-9-1";
    final Set<String> shardIdsOfCurrentLeases = Collections.emptySet();
    final Map<String, Lease> newLeaseMap = Collections.emptyMap();
    final Map<String, Boolean> memoizationContext = new HashMap<>();
    final Map<String, Shard> kinesisShards = new HashMap<>();

    kinesisShards.put(parentShardId, ShardObjectHelper.newShard(parentShardId, null, null, null));
    kinesisShards.put(adjacentParentShardId, ShardObjectHelper.newShard(adjacentParentShardId, null, null, null));
    kinesisShards.put(shardId, ShardObjectHelper.newShard(shardId, parentShardId, adjacentParentShardId, null));

    assertThat(
            HierarchicalShardSyncer.checkIfDescendantAndAddNewLeasesForAncestors(shardId, INITIAL_POSITION_LATEST,
            shardIdsOfCurrentLeases, kinesisShards, newLeaseMap, memoizationContext), equalTo(false));
    assertThat(newLeaseMap.isEmpty(), equalTo(true));
}
 

/**
 * Test method for call().
 * 
 * @throws DependencyException
 * @throws InvalidStateException
 * @throws ProvisionedThroughputException
 */
@Test
public final void testCall() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
    if (!leaseRefresher.leaseTableExists()) {
        final Long readCapacity = 10L;
        final Long writeCapacity = 10L;
        leaseRefresher.createLeaseTableIfNotExists(readCapacity, writeCapacity);
    }
    leaseRefresher.deleteAll();
    Set<String> shardIds = shardDetector.listShards().stream().map(Shard::shardId).collect(Collectors.toSet());
    ShardSyncTask syncTask = new ShardSyncTask(shardDetector, leaseRefresher,
            InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.LATEST), false, false, 0L,
            hierarchicalShardSyncer, NULL_METRICS_FACTORY);
    syncTask.call();
    List<Lease> leases = leaseRefresher.listLeases();
    Set<String> leaseKeys = new HashSet<>();
    for (Lease lease : leases) {
        leaseKeys.add(lease.leaseKey());
    }

    // Verify that all shardIds had leases for them
    Assert.assertEquals(shardIds.size(), leases.size());
    shardIds.removeAll(leaseKeys);
    Assert.assertTrue(shardIds.isEmpty());
}
 

/**
 * Test determineNewLeasesToCreate() where there are no leases and no resharding operations have been performed
 */
@Test
public void testDetermineNewLeasesToCreate0Leases0Reshards() {
    final String shardId0 = "shardId-0";
    final String shardId1 = "shardId-1";
    final SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null);

    final List<Shard> shards = Arrays.asList(ShardObjectHelper.newShard(shardId0, null, null, sequenceRange),
            ShardObjectHelper.newShard(shardId1, null, null, sequenceRange));
    final List<Lease> currentLeases = Collections.emptyList();

    final List<Lease> newLeases = HierarchicalShardSyncer.determineNewLeasesToCreate(shards, currentLeases,
            INITIAL_POSITION_LATEST);
    final Set<String> newLeaseKeys = newLeases.stream().map(Lease::leaseKey).collect(Collectors.toSet());
    final Set<String> expectedLeaseShardIds = new HashSet<>(Arrays.asList(shardId0, shardId1));

    assertThat(newLeases.size(), equalTo(expectedLeaseShardIds.size()));
    assertThat(newLeaseKeys, equalTo(expectedLeaseShardIds));
}
 

public static Shard getShard(final KinesisClient kinesisClient, final String streamName, final String shardIdStart)
		throws Exception {
	LOG.debug(String.format("Getting Shard %s for Stream %s", shardIdStart, streamName));

	KinesisOperation describe = new KinesisOperation() {
		public Object run(KinesisClient client) {
			// reduce the shardIdStart by 1 as the API uses it as an exclusive start key not
			// a filter
			String shardIdToQuery = new BigDecimal(shardIdStart).subtract(new BigDecimal("1")).toString();
			ListShardsRequest req = ListShardsRequest.builder().streamName(streamName)
					.exclusiveStartShardId(shardIdToQuery).build();
			ListShardsResponse result = client.listShards(req);

			return result.shards().get(0);
		}
	};
	return (Shard) doOperation(kinesisClient, describe, streamName, DESCRIBE_RETRIES, false);
}
 

private List<Shard> constructShardListForGraphB() {
    final SequenceNumberRange range0 = ShardObjectHelper.newSequenceNumberRange("1000", "1049");
    final SequenceNumberRange range1 = ShardObjectHelper.newSequenceNumberRange("1050", "1099");
    final SequenceNumberRange range2 = ShardObjectHelper.newSequenceNumberRange("1100", "1149");
    final SequenceNumberRange range3 = ShardObjectHelper.newSequenceNumberRange("1150", "1199");
    final SequenceNumberRange range4 = ShardObjectHelper.newSequenceNumberRange("1200", "1249");
    final SequenceNumberRange range5 = ShardObjectHelper.newSequenceNumberRange("1250", "1299");
    final SequenceNumberRange range6 = ShardObjectHelper.newSequenceNumberRange("1300", null);

    final HashKeyRange hashRange0 = ShardObjectHelper.newHashKeyRange("0", "499");
    final HashKeyRange hashRange1 = ShardObjectHelper.newHashKeyRange("500", ShardObjectHelper.MAX_HASH_KEY);
    final HashKeyRange hashRange2 = ShardObjectHelper.newHashKeyRange("0", ShardObjectHelper.MAX_HASH_KEY);

    return Arrays.asList(ShardObjectHelper.newShard("shardId-0", null, null, range0, hashRange0),
            ShardObjectHelper.newShard("shardId-1", null, null, range0, hashRange1),
            ShardObjectHelper.newShard("shardId-2", "shardId-0", "shardId-1", range1, hashRange2),
            ShardObjectHelper.newShard("shardId-3", "shardId-2", null, range2, hashRange0),
            ShardObjectHelper.newShard("shardId-4", "shardId-2", null, range2, hashRange1),
            ShardObjectHelper.newShard("shardId-5", "shardId-3", "shardId-4", range3, hashRange2),
            ShardObjectHelper.newShard("shardId-6", "shardId-5", null, range4, hashRange0),
            ShardObjectHelper.newShard("shardId-7", "shardId-5", null, range4, hashRange1),
            ShardObjectHelper.newShard("shardId-8", "shardId-6", "shardId-7", range5, hashRange2),
            ShardObjectHelper.newShard("shardId-9", "shardId-8", null, range6, hashRange0),
            ShardObjectHelper.newShard("shardId-10", null, "shardId-8", range6, hashRange1));
}
 

@Test
public void testDetermineNewLeasesToCreateSplitMergeLatest2() {
    final List<Shard> shards = constructShardListForGraphA();
    final List<Lease> currentLeases = Arrays.asList(newLease("shardId-4"), newLease("shardId-5"),
            newLease("shardId-7"));

    final List<Lease> newLeases = HierarchicalShardSyncer.determineNewLeasesToCreate(shards, currentLeases,
            INITIAL_POSITION_LATEST);

    final Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>();
    expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON);
    expectedShardIdCheckpointMap.put("shardId-9", ExtendedSequenceNumber.TRIM_HORIZON);
    expectedShardIdCheckpointMap.put("shardId-10", ExtendedSequenceNumber.TRIM_HORIZON);
    expectedShardIdCheckpointMap.put("shardId-6", ExtendedSequenceNumber.LATEST);

    assertThat(newLeases.size(), equalTo(expectedShardIdCheckpointMap.size()));
    for (Lease lease : newLeases) {
        assertThat("Unexpected lease: " + lease, expectedShardIdCheckpointMap.containsKey(lease.leaseKey()),
                equalTo(true));
        assertThat(lease.checkpoint(), equalTo(expectedShardIdCheckpointMap.get(lease.leaseKey())));
    }
}
 

@Test
public void testDetermineNewLeasesToCreateSplitMergeLatest1() {
    final List<Shard> shards = constructShardListForGraphA();
    final List<Lease> currentLeases = Arrays.asList(newLease("shardId-3"), newLease("shardId-4"),
            newLease("shardId-5"));

    final List<Lease> newLeases = HierarchicalShardSyncer.determineNewLeasesToCreate(shards, currentLeases,
            INITIAL_POSITION_LATEST);

    final Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>();
    expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON);
    expectedShardIdCheckpointMap.put("shardId-9", ExtendedSequenceNumber.TRIM_HORIZON);
    expectedShardIdCheckpointMap.put("shardId-10", ExtendedSequenceNumber.TRIM_HORIZON);
    expectedShardIdCheckpointMap.put("shardId-6", ExtendedSequenceNumber.LATEST);
    expectedShardIdCheckpointMap.put("shardId-2", ExtendedSequenceNumber.LATEST);
    expectedShardIdCheckpointMap.put("shardId-7", ExtendedSequenceNumber.TRIM_HORIZON);

    assertThat(newLeases.size(), equalTo(expectedShardIdCheckpointMap.size()));
    for (Lease lease : newLeases) {
        assertThat("Unexpected lease: " + lease, expectedShardIdCheckpointMap.containsKey(lease.leaseKey()),
                equalTo(true));
        assertThat(lease.checkpoint(), equalTo(expectedShardIdCheckpointMap.get(lease.leaseKey())));
    }
}
 

@Test
public void testDetermineNewLeasesToCreateIgnoreClosedShard() {
    final String lastShardId = "shardId-1";
    final List<Lease> currentLeases = new ArrayList<>();

    final List<Shard> shards = Arrays.asList(
            ShardObjectHelper.newShard("shardId-0", null, null,
                    ShardObjectHelper.newSequenceNumberRange("303", "404")),
            ShardObjectHelper.newShard(lastShardId, null, null,
                    ShardObjectHelper.newSequenceNumberRange("405", null)));

    final List<Lease> newLeases = HierarchicalShardSyncer.determineNewLeasesToCreate(shards, currentLeases,
            INITIAL_POSITION_LATEST);

    assertThat(newLeases.size(), equalTo(1));
    assertThat(newLeases.get(0).leaseKey(), equalTo(lastShardId));
}
 

@Override
@Synchronized
public List<Shard> listShards() {
    final List<Shard> shards = new ArrayList<>();
    ListShardsResponse result;
    String nextToken = null;

    do {
        result = listShards(nextToken);

        if (result == null) {
            /*
             * If listShards ever returns null, we should bail and return null. This indicates the stream is not
             * in ACTIVE or UPDATING state and we may not have accurate/consistent information about the stream.
             */
            return null;
        } else {
            shards.addAll(result.shards());
            nextToken = result.nextToken();
        }
    } while (StringUtils.isNotEmpty(result.nextToken()));

    cachedShardMap(shards);
    return shards;
}
 

/**
 * Test method for {@link ShutdownTask#call()}.
 * This test is for the scenario that checkAndCreateLeaseForNewShards throws an exception.
 */
@Test
public final void testCallWhenSyncingShardsThrows() throws Exception {
    List<Shard> latestShards = constructShardListGraphA();
    when(shardDetector.listShards()).thenReturn(latestShards);
    when(recordProcessorCheckpointer.lastCheckpointValue()).thenReturn(ExtendedSequenceNumber.SHARD_END);
    when(leaseCoordinator.leaseRefresher()).thenReturn(leaseRefresher);

    doAnswer((invocation) -> {
        throw new KinesisClientLibIOException("KinesisClientLibIOException");
    }).when(hierarchicalShardSyncer)
            .checkAndCreateLeaseForNewShards(shardDetector, leaseRefresher, INITIAL_POSITION_TRIM_HORIZON,
                    cleanupLeasesOfCompletedShards, ignoreUnexpectedChildShards,
                    NULL_METRICS_FACTORY.createMetrics(), latestShards);

    final TaskResult result = task.call();
    assertNotNull(result.getException());
    assertTrue(result.getException() instanceof KinesisClientLibIOException);
    verify(recordsPublisher).shutdown();
    verify(shardRecordProcessor).shardEnded(ShardEndedInput.builder().checkpointer(recordProcessorCheckpointer).build());
    verify(shardRecordProcessor, never()).leaseLost(LeaseLostInput.builder().build());
}
 

/**
 * Compares two leases based on the starting sequence number of corresponding shards.
 * If shards are not found in the shardId->shard map supplied, we do a string comparison on the shardIds.
 * We assume that lease1 and lease2 are:
 *     a/ not null,
 *     b/ shards (if found) have non-null starting sequence numbers
 * 
 * {@inheritDoc}
 */
@Override
public int compare(final Lease lease1, final Lease lease2) {
    int result = 0;
    final String shardId1 = lease1.leaseKey();
    final String shardId2 = lease2.leaseKey();
    final Shard shard1 = shardIdToShardMap.get(shardId1);
    final Shard shard2 = shardIdToShardMap.get(shardId2);
    
    // If we found shards for the two leases, use comparison of the starting sequence numbers
    if (shard1 != null && shard2 != null) {
        BigInteger sequenceNumber1 = new BigInteger(shard1.sequenceNumberRange().startingSequenceNumber());
        BigInteger sequenceNumber2 = new BigInteger(shard2.sequenceNumberRange().startingSequenceNumber());
        result = sequenceNumber1.compareTo(sequenceNumber2);                
    }
    
    if (result == 0) {
        result = shardId1.compareTo(shardId2);
    }
    
    return result;
}
 

@Test
public void testGetShardNonExistentShardForceRefresh() {
    final String shardId = String.format(SHARD_ID, 5);
    final CompletableFuture<ListShardsResponse> future = CompletableFuture
            .completedFuture(ListShardsResponse.builder().shards(createShardList()).build());

    shardDetector.cachedShardMap(createShardList());

    when(client.listShards(any(ListShardsRequest.class))).thenReturn(future);

    final List<Shard> responses = IntStream.range(0, MAX_CACHE_MISSES_BEFORE_RELOAD + 1)
            .mapToObj(x -> shardDetector.shard(shardId)).collect(Collectors.toList());

    responses.forEach(response -> assertThat(response, nullValue()));
    assertThat(shardDetector.cacheMisses().get(), equalTo(0));
    verify(client).listShards(any(ListShardsRequest.class));
}
 

@Test
public void testGetShardNewShardForceRefresh() {
    final String shardId = String.format(SHARD_ID, 5);
    final List<Shard> shards = new ArrayList<>(createShardList());
    shards.add(Shard.builder().shardId(shardId).build());

    final CompletableFuture<ListShardsResponse> future = CompletableFuture
            .completedFuture(ListShardsResponse.builder().shards(shards).build());

    shardDetector.cachedShardMap(createShardList());

    when(client.listShards(any(ListShardsRequest.class))).thenReturn(future);

    final List<Shard> responses = IntStream.range(0, MAX_CACHE_MISSES_BEFORE_RELOAD + 1)
            .mapToObj(x -> shardDetector.shard(shardId)).collect(Collectors.toList());

    IntStream.range(0, MAX_CACHE_MISSES_BEFORE_RELOAD).forEach(x -> {
        assertThat(responses.get(x), nullValue());
    });

    assertThat(responses.get(MAX_CACHE_MISSES_BEFORE_RELOAD), equalTo(Shard.builder().shardId(shardId).build()));
    verify(client).listShards(any(ListShardsRequest.class));
}
 

@Test
public void testGetShardEmptyCache() {
    final String shardId = String.format(SHARD_ID, 1);
    final CompletableFuture<ListShardsResponse> future = CompletableFuture
            .completedFuture(ListShardsResponse.builder().shards(createShardList()).build());

    when(client.listShards(any(ListShardsRequest.class))).thenReturn(future);

    final Shard shard = shardDetector.shard(shardId);

    assertThat(shard, equalTo(Shard.builder().shardId(shardId).build()));
    verify(client).listShards(any(ListShardsRequest.class));
}
 

@Test
public void testGetShardNonExistentShard() {
    final String shardId = String.format(SHARD_ID, 5);

    shardDetector.cachedShardMap(createShardList());

    final Shard shard = shardDetector.shard(shardId);

    assertThat(shard, nullValue());
    assertThat(shardDetector.cacheMisses().get(), equalTo(1));
    verify(client, never()).listShards(any(ListShardsRequest.class));
}
 

/**
 * Test checkAndCreateLeaseForNewShards with an empty list of shards. In this scenario, shardDetector.listShards()
 * should never be called.
 */
@Test
public void testCheckAndCreateLeasesForShardsWithEmptyShardList() throws Exception {
    final List<Shard> shards = constructShardListForGraphA();

    final ArgumentCaptor<Lease> leaseCaptor = ArgumentCaptor.forClass(Lease.class);
    when(shardDetector.listShards()).thenReturn(shards);
    when(dynamoDBLeaseRefresher.listLeases()).thenReturn(Collections.emptyList());
    when(dynamoDBLeaseRefresher.createLeaseIfNotExists(leaseCaptor.capture())).thenReturn(true);

    hierarchicalShardSyncer
            .checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
                                             cleanupLeasesOfCompletedShards, false, SCOPE, new ArrayList<Shard>());

    final Set<String> expectedShardIds = new HashSet<>();

    final List<Lease> requestLeases = leaseCaptor.getAllValues();
    final Set<String> requestLeaseKeys = requestLeases.stream().map(Lease::leaseKey).collect(Collectors.toSet());
    final Set<ExtendedSequenceNumber> extendedSequenceNumbers = requestLeases.stream().map(Lease::checkpoint)
                                                                             .collect(Collectors.toSet());

    assertThat(requestLeases.size(), equalTo(expectedShardIds.size()));
    assertThat(extendedSequenceNumbers.size(), equalTo(0));

    verify(shardDetector, never()).listShards();
    verify(dynamoDBLeaseRefresher, never()).createLeaseIfNotExists(any(Lease.class));
    verify(dynamoDBLeaseRefresher, never()).deleteLease(any(Lease.class));
}
 

/**
 * Test checkAndCreateLeaseForNewShards with a pre-fetched list of shards. In this scenario, shardDetector.listShards()
 * should never be called.
 */
@Test
public void testCheckAndCreateLeasesForShardsWithShardList() throws Exception {
    final List<Shard> latestShards = constructShardListForGraphA();

    final ArgumentCaptor<Lease> leaseCaptor = ArgumentCaptor.forClass(Lease.class);
    when(shardDetector.listShards()).thenReturn(latestShards);
    when(dynamoDBLeaseRefresher.listLeases()).thenReturn(Collections.emptyList());
    when(dynamoDBLeaseRefresher.createLeaseIfNotExists(leaseCaptor.capture())).thenReturn(true);

    hierarchicalShardSyncer
            .checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
                                             cleanupLeasesOfCompletedShards, false, SCOPE, latestShards);

    final Set<String> expectedShardIds = new HashSet<>(
            Arrays.asList("shardId-4", "shardId-8", "shardId-9", "shardId-10"));

    final List<Lease> requestLeases = leaseCaptor.getAllValues();
    final Set<String> requestLeaseKeys = requestLeases.stream().map(Lease::leaseKey).collect(Collectors.toSet());
    final Set<ExtendedSequenceNumber> extendedSequenceNumbers = requestLeases.stream().map(Lease::checkpoint)
                                                                             .collect(Collectors.toSet());

    assertThat(requestLeases.size(), equalTo(expectedShardIds.size()));
    assertThat(requestLeaseKeys, equalTo(expectedShardIds));
    assertThat(extendedSequenceNumbers.size(), equalTo(1));

    extendedSequenceNumbers.forEach(seq -> assertThat(seq, equalTo(ExtendedSequenceNumber.LATEST)));

    verify(shardDetector, never()).listShards();
    verify(dynamoDBLeaseRefresher, times(expectedShardIds.size())).createLeaseIfNotExists(any(Lease.class));
    verify(dynamoDBLeaseRefresher, never()).deleteLease(any(Lease.class));
}
 

/**
 * Test checkAndCreateLeaseForNewShards while not providing a pre-fetched list of shards
 */
@Test
public void testCheckAndCreateLeasesForShardsIfMissingAtLatest() throws Exception {
    final List<Shard> shards = constructShardListForGraphA();

    final ArgumentCaptor<Lease> leaseCaptor = ArgumentCaptor.forClass(Lease.class);

    when(shardDetector.listShards()).thenReturn(shards);
    when(dynamoDBLeaseRefresher.listLeases()).thenReturn(Collections.emptyList());
    when(dynamoDBLeaseRefresher.createLeaseIfNotExists(leaseCaptor.capture())).thenReturn(true);

    hierarchicalShardSyncer
            .checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
            cleanupLeasesOfCompletedShards, false, SCOPE);

    final Set<String> expectedShardIds = new HashSet<>(
            Arrays.asList("shardId-4", "shardId-8", "shardId-9", "shardId-10"));

    final List<Lease> requestLeases = leaseCaptor.getAllValues();
    final Set<String> requestLeaseKeys = requestLeases.stream().map(Lease::leaseKey).collect(Collectors.toSet());
    final Set<ExtendedSequenceNumber> extendedSequenceNumbers = requestLeases.stream().map(Lease::checkpoint)
            .collect(Collectors.toSet());

    assertThat(requestLeases.size(), equalTo(expectedShardIds.size()));
    assertThat(requestLeaseKeys, equalTo(expectedShardIds));
    assertThat(extendedSequenceNumbers.size(), equalTo(1));

    extendedSequenceNumbers.forEach(seq -> assertThat(seq, equalTo(ExtendedSequenceNumber.LATEST)));

    verify(shardDetector).listShards();
    verify(dynamoDBLeaseRefresher, times(expectedShardIds.size())).createLeaseIfNotExists(any(Lease.class));
    verify(dynamoDBLeaseRefresher, never()).deleteLease(any(Lease.class));

}
 

/**
 * Helper method to get parent shardIds of the current shard - includes the parent shardIds if:
 * a/ they are not null
 * b/ if they exist in the current shard map (i.e. haven't expired)
 * 
 * @param shard Will return parents of this shard
 * @param shardIdToShardMapOfAllKinesisShards ShardId->Shard map containing all shards obtained via DescribeStream.
 * @return Set of parentShardIds
 */
static Set<String> getParentShardIds(final Shard shard,
        final Map<String, Shard> shardIdToShardMapOfAllKinesisShards) {
    final Set<String> parentShardIds = new HashSet<>(2);
    final String parentShardId = shard.parentShardId();
    if (parentShardId != null && shardIdToShardMapOfAllKinesisShards.containsKey(parentShardId)) {
        parentShardIds.add(parentShardId);
    }
    final String adjacentParentShardId = shard.adjacentParentShardId();
    if (adjacentParentShardId != null && shardIdToShardMapOfAllKinesisShards.containsKey(adjacentParentShardId)) {
        parentShardIds.add(adjacentParentShardId);
    }
    return parentShardIds;
}
 

static List<String> getParentShardIds(Shard shard) {
    List<String> parentShardIds = new ArrayList<>(2);
    if (shard.adjacentParentShardId() != null) {
        parentShardIds.add(shard.adjacentParentShardId());
    }
    if (shard.parentShardId() != null) {
        parentShardIds.add(shard.parentShardId());
    }
    return parentShardIds;
}
 

private List<Shard> constructShardListGraphA() {
    final SequenceNumberRange range0 = ShardObjectHelper.newSequenceNumberRange("11", "102");
    final SequenceNumberRange range1 = ShardObjectHelper.newSequenceNumberRange("11", null);
    final SequenceNumberRange range2 = ShardObjectHelper.newSequenceNumberRange("11", "205");
    final SequenceNumberRange range3 = ShardObjectHelper.newSequenceNumberRange("103", "205");
    final SequenceNumberRange range4 = ShardObjectHelper.newSequenceNumberRange("206", null);

    return Arrays.asList(
            ShardObjectHelper.newShard("shardId-0", null, null, range0,
                                       ShardObjectHelper.newHashKeyRange("0", "99")),
            ShardObjectHelper.newShard("shardId-1", null, null, range0,
                                       ShardObjectHelper.newHashKeyRange("100", "199")),
            ShardObjectHelper.newShard("shardId-2", null, null, range0,
                                       ShardObjectHelper.newHashKeyRange("200", "299")),
            ShardObjectHelper.newShard("shardId-3", null, null, range0,
                                       ShardObjectHelper.newHashKeyRange("300", "399")),
            ShardObjectHelper.newShard("shardId-4", null, null, range1,
                                       ShardObjectHelper.newHashKeyRange("400", "499")),
            ShardObjectHelper.newShard("shardId-5", null, null, range2,
                                       ShardObjectHelper.newHashKeyRange("500", ShardObjectHelper.MAX_HASH_KEY)),
            ShardObjectHelper.newShard("shardId-6", "shardId-0", "shardId-1", range3,
                                       ShardObjectHelper.newHashKeyRange("0", "199")),
            ShardObjectHelper.newShard("shardId-7", "shardId-2", "shardId-3", range3,
                                       ShardObjectHelper.newHashKeyRange("200", "399")),
            ShardObjectHelper.newShard("shardId-8", "shardId-6", "shardId-7", range4,
                                       ShardObjectHelper.newHashKeyRange("0", "399")),
            ShardObjectHelper.newShard("shardId-9", "shardId-5", null, range4,
                                       ShardObjectHelper.newHashKeyRange("500", "799")),
            ShardObjectHelper.newShard("shardId-10", null, "shardId-5", range4,
                                       ShardObjectHelper.newHashKeyRange("800", ShardObjectHelper.MAX_HASH_KEY)));
}
 

private List<Shard> createShardList() {
    return Arrays.asList(Shard.builder().shardId(String.format(SHARD_ID, 0)).build(),
            Shard.builder().shardId(String.format(SHARD_ID, 1)).build(),
            Shard.builder().shardId(String.format(SHARD_ID, 2)).build(),
            Shard.builder().shardId(String.format(SHARD_ID, 3)).build(),
            Shard.builder().shardId(String.format(SHARD_ID, 4)).build());
}
 

private List<Lease> createLeasesFromShards(final List<Shard> shards, final ExtendedSequenceNumber checkpoint,
        final String leaseOwner) {
    return shards.stream().map(shard -> {
        final Set<String> parentShardIds = new HashSet<>();
        if (StringUtils.isNotEmpty(shard.parentShardId())) {
            parentShardIds.add(shard.parentShardId());
        }
        if (StringUtils.isNotEmpty(shard.adjacentParentShardId())) {
            parentShardIds.add(shard.adjacentParentShardId());
        }
        return new Lease(shard.shardId(), leaseOwner, 0L, UUID.randomUUID(), 0L, checkpoint, null, 0L,
                parentShardIds);
    }).collect(Collectors.toList());
}
 

@Test
public void testCleanUpGarbageLeaseForNonExistentShard() throws Exception {
    final List<Shard> shards = constructShardListForGraphA();
    final String garbageShardId = "shardId-garbage-001";
    final Shard garbageShard = ShardObjectHelper.newShard(garbageShardId, null, null,
            ShardObjectHelper.newSequenceNumberRange("101", null));
    final Lease garbageLease = createLeaseFromShard(garbageShard, new ExtendedSequenceNumber("99"), LEASE_OWNER);
    final List<Lease> leases = new ArrayList<>(
            createLeasesFromShards(shards, ExtendedSequenceNumber.TRIM_HORIZON, LEASE_OWNER));
    leases.add(garbageLease);

    final ArgumentCaptor<Lease> leaseCaptor = ArgumentCaptor.forClass(Lease.class);

    when(shardDetector.listShards()).thenReturn(shards);
    when(dynamoDBLeaseRefresher.listLeases()).thenReturn(leases);
    doNothing().when(dynamoDBLeaseRefresher).deleteLease(leaseCaptor.capture());

    hierarchicalShardSyncer.checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher,
            INITIAL_POSITION_TRIM_HORIZON, cleanupLeasesOfCompletedShards, ignoreUnexpectedChildShards, SCOPE);

    assertThat(leaseCaptor.getAllValues().size(), equalTo(1));
    assertThat(leaseCaptor.getValue(), equalTo(garbageLease));

    verify(shardDetector, times(2)).listShards();
    verify(dynamoDBLeaseRefresher).listLeases();
    verify(dynamoDBLeaseRefresher).deleteLease(any(Lease.class));
    verify(dynamoDBLeaseRefresher, never()).createLeaseIfNotExists(any(Lease.class));
}
 

/**
 * Test determineNewLeasesToCreate() where there are no shards
 */
@Test
public void testDetermineNewLeasesToCreateNoShards() {
    final List<Shard> shards = Collections.emptyList();
    final List<Lease> leases = Collections.emptyList();

    assertThat(HierarchicalShardSyncer.determineNewLeasesToCreate(shards, leases, INITIAL_POSITION_LATEST).isEmpty(),
            equalTo(true));
}
 

@Test
public void testListShardsResouceInUse() {
    final CompletableFuture<ListShardsResponse> future = CompletableFuture.supplyAsync(() -> {
        throw ResourceInUseException.builder().build();
    });

    when(client.listShards(any(ListShardsRequest.class))).thenReturn(future);

    final List<Shard> shards = shardDetector.listShards();

    assertThat(shards, nullValue());
    verify(client).listShards(any(ListShardsRequest.class));

}