Fix the unit tests

Signed-off-by: Souvik Bose <[email protected]>
opensearch-project · Dec 13, 2024 · c30dc2d · c30dc2d
1 parent 634c031
commit c30dc2d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/.../opensearch/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbClientWrapperTest.java b/.../opensearch/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbClientWrapperTest.java
@@ -62,7 +62,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
@@ -518,6 +518,8 @@ void getAvailablePartition_will_continue_until_tryAcquirePartition_succeeds(fina
 
         final Duration ttl = Duration.ofSeconds(new Random().nextInt());
 
+        final Long expectedExpirationTime = Instant.now().plus(ttl).getEpochSecond();
+
         final Optional<SourcePartitionStoreItem> result = objectUnderTest.getAvailablePartition(
                 ownerId, ownershipTimeout, SourcePartitionStatus.valueOf(sourcePartitionStatus), sourceStatusCombinationKey, new Random().nextInt(20), ttl);
 
@@ -538,7 +540,7 @@ void getAvailablePartition_will_continue_until_tryAcquirePartition_succeeds(fina
 
         final ArgumentCaptor<Long> expiryTimeArgumentCaptor = ArgumentCaptor.forClass(Long.class);
         verify(acquiredItem).setExpirationTime(expiryTimeArgumentCaptor.capture());
-        assertThat(expiryTimeArgumentCaptor.getValue(), greaterThan(Instant.now().getEpochSecond()));
+        assertThat(expiryTimeArgumentCaptor.getValue(), greaterThanOrEqualTo(expectedExpirationTime));
 
         verify(acquiredItem).setPartitionPriority(newPartitionOwnershipTimeout.toString());
     }
@@ -584,6 +586,7 @@ void getAvailablePartition_with_multiple_pages_continue_until_tryAcquirePartitio
         reflectivelySetField(objectUnderTest, "table", table);
 
         final Duration ttl = Duration.ofSeconds(new Random().nextInt());
+        final Long expectedExpirationTime = Instant.now().plus(ttl).getEpochSecond();
 
         final Optional<SourcePartitionStoreItem> result = objectUnderTest.getAvailablePartition(
                 ownerId, ownershipTimeout, SourcePartitionStatus.valueOf(sourcePartitionStatus), sourceStatusCombinationKey, new Random().nextInt(20), ttl);
@@ -607,7 +610,7 @@ void getAvailablePartition_with_multiple_pages_continue_until_tryAcquirePartitio
 
         final ArgumentCaptor<Long> expiryTimeArgumentCaptor = ArgumentCaptor.forClass(Long.class);
         verify(acquiredItem).setExpirationTime(expiryTimeArgumentCaptor.capture());
-        assertThat(expiryTimeArgumentCaptor.getValue(), greaterThan(Instant.now().getEpochSecond()));
+        assertThat(expiryTimeArgumentCaptor.getValue(), greaterThanOrEqualTo(expectedExpirationTime));
     }
 
     @ParameterizedTest
@@ -651,6 +654,7 @@ void getAvailablePartition_with_multiple_pages_will_iterate_through_all_items_wi
         reflectivelySetField(objectUnderTest, "table", table);
 
         final Duration ttl = Duration.ofSeconds(new Random().nextInt());
+        final Long expectedExpirationTime = Instant.now().plus(ttl).getEpochSecond();
 
         final Optional<SourcePartitionStoreItem> result = objectUnderTest.getAvailablePartition(
                 ownerId, ownershipTimeout, SourcePartitionStatus.valueOf(sourcePartitionStatus), sourceStatusCombinationKey, new Random().nextInt(20), ttl);
@@ -673,7 +677,7 @@ void getAvailablePartition_with_multiple_pages_will_iterate_through_all_items_wi
 
         final ArgumentCaptor<Long> expiryTimeArgumentCaptor = ArgumentCaptor.forClass(Long.class);
         verify(acquiredItem).setExpirationTime(expiryTimeArgumentCaptor.capture());
-        assertThat(expiryTimeArgumentCaptor.getValue(), greaterThan(Instant.now().getEpochSecond()));
+        assertThat(expiryTimeArgumentCaptor.getValue(), greaterThanOrEqualTo(expectedExpirationTime));
     }
 
     @Test

diff --git a/...h/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbSourceCoordinationStoreTest.java b/...h/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbSourceCoordinationStoreTest.java
@@ -247,6 +247,8 @@ void getAvailablePartition_with_no_item_acquired_returns_empty_optional() {
         final Duration ownershipTimeout = Duration.ofMinutes(2);
         final Duration ttl = Duration.ofSeconds(new Random().nextInt());
 
+        given(dynamoStoreSettings.getTtl()).willReturn(ttl);
+
         given(dynamoDbClientWrapper.getAvailablePartition(ownerId, ownershipTimeout,
                 SourcePartitionStatus.ASSIGNED,
                 String.format(SOURCE_STATUS_COMBINATION_KEY_FORMAT, sourceIdentifier, SourcePartitionStatus.ASSIGNED),
@@ -275,6 +277,7 @@ void getAvailablePartition_with_acquired_ASSIGNED_partition_returns_the_partitio
         final Duration ownershipTimeout = Duration.ofMinutes(2);
         final Duration ttl = Duration.ofSeconds(new Random().nextInt());
 
+        given(dynamoStoreSettings.getTtl()).willReturn(ttl);
         final DynamoDbSourcePartitionItem acquiredItem = mock(DynamoDbSourcePartitionItem.class);
 
         given(dynamoDbClientWrapper.getAvailablePartition(ownerId, ownershipTimeout,
@@ -298,6 +301,7 @@ void getAvailablePartition_with_acquired_CLOSED_partition_returns_the_partition(
         final Duration ownershipTimeout = Duration.ofMinutes(2);
         final Duration ttl = Duration.ofSeconds(new Random().nextInt());
 
+        given(dynamoStoreSettings.getTtl()).willReturn(ttl);
         final DynamoDbSourcePartitionItem acquiredItem = mock(DynamoDbSourcePartitionItem.class);
 
         given(dynamoDbClientWrapper.getAvailablePartition(ownerId, ownershipTimeout,
@@ -330,6 +334,7 @@ void getAvailablePartition_with_acquired_UNASSIGNED_partition_returns_the_partit
         final String sourceIdentifier = UUID.randomUUID().toString();
         final Duration ownershipTimeout = Duration.ofMinutes(2);
         final Duration ttl = Duration.ofSeconds(new Random().nextInt());
+        given(dynamoStoreSettings.getTtl()).willReturn(ttl);
 
         final DynamoDbSourcePartitionItem acquiredItem = mock(DynamoDbSourcePartitionItem.class);