|
27 | 27 | import org.opensearch.dataprepper.plugins.source.ownership.BucketOwnerProvider; |
28 | 28 | import software.amazon.awssdk.services.s3.S3Client; |
29 | 29 | import software.amazon.awssdk.services.s3.model.DeleteObjectRequest; |
| 30 | +import software.amazon.awssdk.services.s3.model.NoSuchKeyException; |
30 | 31 |
|
31 | 32 | import java.io.IOException; |
32 | 33 | import java.time.Duration; |
|
47 | 48 | import static org.mockito.Mockito.doAnswer; |
48 | 49 | import static org.mockito.Mockito.doNothing; |
49 | 50 | import static org.mockito.Mockito.doThrow; |
50 | | -import static org.mockito.Mockito.when; |
51 | 51 | import static org.mockito.Mockito.verify; |
52 | 52 | import static org.mockito.Mockito.verifyNoInteractions; |
53 | 53 | import static org.mockito.Mockito.verifyNoMoreInteractions; |
| 54 | +import static org.mockito.Mockito.when; |
54 | 55 | import static org.opensearch.dataprepper.plugins.source.ScanObjectWorker.ACKNOWLEDGEMENT_SET_CALLBACK_METRIC_NAME; |
55 | 56 |
|
56 | 57 | @ExtendWith(MockitoExtension.class) |
@@ -275,6 +276,26 @@ void getNextPartition_supplier_is_expected_partitionCreationSupplier() { |
275 | 276 | objectUnderTest.runWithoutInfiniteLoop(); |
276 | 277 | } |
277 | 278 |
|
| 279 | + @Test |
| 280 | + void partitionIsCompleted_when_NoObjectKeyException_is_thrown_from_process_object() throws IOException { |
| 281 | + final String bucket = UUID.randomUUID().toString(); |
| 282 | + final String objectKey = UUID.randomUUID().toString(); |
| 283 | + final String partitionKey = bucket + "|" + objectKey; |
| 284 | + |
| 285 | + |
| 286 | + final SourcePartition<S3SourceProgressState> partitionToProcess = SourcePartition.builder(S3SourceProgressState.class).withPartitionKey(partitionKey).build(); |
| 287 | + |
| 288 | + given(sourceCoordinator.getNextPartition(any(Function.class))).willReturn(Optional.of(partitionToProcess)); |
| 289 | + |
| 290 | + final ArgumentCaptor<S3ObjectReference> objectReferenceArgumentCaptor = ArgumentCaptor.forClass(S3ObjectReference.class); |
| 291 | + doThrow(NoSuchKeyException.class).when(s3ObjectHandler).parseS3Object(objectReferenceArgumentCaptor.capture(), eq(null), eq(sourceCoordinator), eq(partitionKey)); |
| 292 | + doNothing().when(sourceCoordinator).completePartition(partitionKey); |
| 293 | + |
| 294 | + createObjectUnderTest().runWithoutInfiniteLoop(); |
| 295 | + |
| 296 | + verifyNoMoreInteractions(sourceCoordinator); |
| 297 | + } |
| 298 | + |
278 | 299 | static Stream<Class> exceptionProvider() { |
279 | 300 | return Stream.of(PartitionUpdateException.class, PartitionNotFoundException.class, PartitionNotOwnedException.class); |
280 | 301 | } |
|
0 commit comments