One of our new DDB integration tests is failing on GitHub. This change attempts to fix that by 1) including a sleep between writing events to ensure that they have different timestamps; and 2) waiting for the GSI to reach eventual consistency.

Signed-off-by: David Venable <dlv@amazon.com>

Author: dlvenable

data-prepper-plugins/dynamodb-source-coordination-store/src/test/java/org/opensearch/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbSourceCoordinationStoreIT.java

@@ -24,19 +24,30 @@
 import org.opensearch.dataprepper.model.source.coordinator.SourcePartitionStatus;
 import org.opensearch.dataprepper.model.source.coordinator.SourcePartitionStoreItem;
 import org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionUpdateException;
+import software.amazon.awssdk.core.pagination.sync.SdkIterable;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
+import software.amazon.awssdk.enhanced.dynamodb.DynamoDbIndex;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbTable;
 import software.amazon.awssdk.enhanced.dynamodb.Key;
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
 import software.amazon.awssdk.enhanced.dynamodb.model.GetItemEnhancedRequest;
+import software.amazon.awssdk.enhanced.dynamodb.model.Page;
+import software.amazon.awssdk.enhanced.dynamodb.model.QueryConditional;
+import software.amazon.awssdk.enhanced.dynamodb.model.QueryEnhancedRequest;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.List;
 import java.util.Optional;
 import java.util.Random;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
+import static org.awaitility.Awaitility.await;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
@@ -48,6 +59,8 @@
 import static org.hamcrest.Matchers.sameInstance;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.when;
+import static org.opensearch.dataprepper.plugins.sourcecoordinator.dynamodb.DynamoDbClientWrapper.SOURCE_STATUS_COMBINATION_KEY_GLOBAL_SECONDARY_INDEX;
+import static org.opensearch.dataprepper.plugins.sourcecoordinator.dynamodb.DynamoDbSourceCoordinationStore.SOURCE_STATUS_COMBINATION_KEY_FORMAT;
 
 @ExtendWith(MockitoExtension.class)
 class DynamoDbSourceCoordinationStoreIT {
@@ -91,8 +104,8 @@ void setUp() {
         when(dynamoStoreSettings.getRegion()).thenReturn(region);
         when(dynamoStoreSettings.getStsRoleArn()).thenReturn(stsRoleArn);
         when(dynamoStoreSettings.getStsExternalId()).thenReturn(stsExternalId);
-        when(dynamoStoreSettings.getProvisionedReadCapacityUnits()).thenReturn(1L);
-        when(dynamoStoreSettings.getProvisionedWriteCapacityUnits()).thenReturn(1L);
+        when(dynamoStoreSettings.getProvisionedReadCapacityUnits()).thenReturn(10L);
+        when(dynamoStoreSettings.getProvisionedWriteCapacityUnits()).thenReturn(10L);
 
         when(dynamoDbClientFactory.provideDynamoDbClient(region, stsRoleArn, stsExternalId)).thenReturn(dynamoDbClient);
 
@@ -250,36 +263,56 @@ void tryCreatePartitionItem_creates_an_item() {
     }
 
     @Test
-    void tryAcquireAvailablePartition_gets_first_unassigned_partition() {
+    void tryAcquireAvailablePartition_gets_first_unassigned_partition() throws InterruptedException {
         final DynamoDbSourceCoordinationStore objectUnderTest = createObjectUnderTest();
         final String partitionProgressState = UUID.randomUUID().toString();
 
-        final String unassignedPartitionKey1 = UUID.randomUUID().toString();
-        final String unassignedPartitionKey2 = UUID.randomUUID().toString();
-        final String unassignedPartitionKey3 = UUID.randomUUID().toString();
+        final List<String> partitionKeys = IntStream.rangeClosed(1, 3)
+                .mapToObj(i -> UUID.randomUUID() + "_" + i)
+                .collect(Collectors.toList());
 
-        objectUnderTest.tryCreatePartitionItem(sourceIdentifier,
-                unassignedPartitionKey1, SourcePartitionStatus.UNASSIGNED, 1L, partitionProgressState, false);
-        objectUnderTest.tryCreatePartitionItem(sourceIdentifier,
-                unassignedPartitionKey2, SourcePartitionStatus.UNASSIGNED, 1L, partitionProgressState, false);
-        objectUnderTest.tryCreatePartitionItem(sourceIdentifier,
-                unassignedPartitionKey3, SourcePartitionStatus.UNASSIGNED, 1L, partitionProgressState, false);
+        for (final String partitionKey : partitionKeys) {
+            final boolean createSuccess = objectUnderTest.tryCreatePartitionItem(sourceIdentifier,
+                    partitionKey, SourcePartitionStatus.UNASSIGNED, 1L, partitionProgressState, false);
+            assertThat(createSuccess, equalTo(true));
+            Thread.sleep(150);
+        }
 
-        final Optional<SourcePartitionStoreItem> maybeAcquired =
-                objectUnderTest.tryAcquireAvailablePartition(sourceIdentifier, ownerId, Duration.ofSeconds(20));
+        await().atMost(10, TimeUnit.SECONDS).until(() -> {
+            final String primaryKey = String.format(SOURCE_STATUS_COMBINATION_KEY_FORMAT, sourceIdentifier, SourcePartitionStatus.UNASSIGNED);
+            final Stream<DynamoDbSourcePartitionItem> items = querySourceStatusIndex(primaryKey);
+
+            return items.count() == partitionKeys.size();
+        });
+
+        final Optional<SourcePartitionStoreItem> maybeAcquired = objectUnderTest.tryAcquireAvailablePartition(sourceIdentifier, ownerId, Duration.ofSeconds(20));
 
         assertThat(maybeAcquired, notNullValue());
         assertThat(maybeAcquired.isPresent(), equalTo(true));
-
         final SourcePartitionStoreItem acquiredItem = maybeAcquired.get();
 
         assertThat(acquiredItem, notNullValue());
+        final String unassignedPartitionKey1 = partitionKeys.get(0);
+
         assertThat(acquiredItem.getSourceIdentifier(), equalTo(sourceIdentifier));
         assertThat(acquiredItem.getSourcePartitionKey(), equalTo(unassignedPartitionKey1));
         assertThat(acquiredItem.getSourcePartitionStatus(), equalTo(SourcePartitionStatus.ASSIGNED));
         assertThat(acquiredItem.getPartitionOwner(), equalTo(ownerId));
     }
 
+    private Stream<DynamoDbSourcePartitionItem> querySourceStatusIndex(final String partitionKey) {
+        final DynamoDbIndex<DynamoDbSourcePartitionItem> sourceStatusIndex = table.index(SOURCE_STATUS_COMBINATION_KEY_GLOBAL_SECONDARY_INDEX);
+        final QueryEnhancedRequest queryEnhancedRequest = QueryEnhancedRequest.builder()
+                .limit(1)
+                .queryConditional(QueryConditional.keyEqualTo(Key.builder().partitionValue(partitionKey).build()))
+                .build();
+
+        final SdkIterable<Page<DynamoDbSourcePartitionItem>> pages = sourceStatusIndex.query(queryEnhancedRequest);
+
+        return pages.stream()
+                .flatMap(page -> page.items().stream());
+    }
+
     private DynamoDbSourcePartitionItem putDynamoDbSourcePartitionItem(final SourcePartitionStatus sourcePartitionStatus) {
         final DynamoDbSourcePartitionItem putItem = createUnsavedPartitionItem(sourcePartitionStatus);
         table.putItem(putItem);