-
Notifications
You must be signed in to change notification settings - Fork 194
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Make s3 partition size configurable and add unit test for S3 partitio…
…n creator classes (#4437) * Make s3 partition size configurable and add unit test for S3 partition creator classes Signed-off-by: Dinu John <86094133+dinujoh@users.noreply.github.com> * Rename export partition size to export batch size Signed-off-by: Dinu John <86094133+dinujoh@users.noreply.github.com> --------- Signed-off-by: Dinu John <86094133+dinujoh@users.noreply.github.com>
- Loading branch information
Showing
9 changed files
with
189 additions
and
48 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
50 changes: 50 additions & 0 deletions
50
...rg/opensearch/dataprepper/plugins/mongo/s3partition/S3FolderPartitionCoordinatorTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| package org.opensearch.dataprepper.plugins.mongo.s3partition; | ||
|
|
||
| import org.junit.jupiter.api.Test; | ||
| import org.junit.jupiter.api.extension.ExtendWith; | ||
| import org.mockito.InjectMocks; | ||
| import org.mockito.Mock; | ||
| import org.mockito.junit.jupiter.MockitoExtension; | ||
| import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator; | ||
| import org.opensearch.dataprepper.plugins.mongo.coordination.partition.GlobalState; | ||
| import org.opensearch.dataprepper.plugins.mongo.model.S3PartitionStatus; | ||
|
|
||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.Optional; | ||
| import java.util.UUID; | ||
|
|
||
| import static org.hamcrest.CoreMatchers.is; | ||
| import static org.hamcrest.MatcherAssert.assertThat; | ||
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.when; | ||
|
|
||
| @ExtendWith(MockitoExtension.class) | ||
| public class S3FolderPartitionCoordinatorTest { | ||
| @Mock | ||
| private EnhancedSourceCoordinator sourceCoordinator; | ||
|
|
||
| @InjectMocks | ||
| private S3FolderPartitionCoordinator s3FolderPartitionCoordinator; | ||
|
|
||
| @Test | ||
| public void getGlobalS3FolderCreationStatus_empty() { | ||
| final String collection = UUID.randomUUID().toString(); | ||
| when(sourceCoordinator.getPartition(S3PartitionCreatorScheduler.S3_FOLDER_PREFIX + collection)).thenReturn(Optional.empty()); | ||
| Optional<S3PartitionStatus> partitionStatus = s3FolderPartitionCoordinator.getGlobalS3FolderCreationStatus(collection); | ||
| assertThat(partitionStatus.isEmpty(), is(true)); | ||
| } | ||
|
|
||
| @Test | ||
| public void getGlobalS3FolderCreationStatus_nonEmpty() { | ||
| final String collection = UUID.randomUUID().toString(); | ||
| final List<String> partitions = List.of(UUID.randomUUID().toString(), UUID.randomUUID().toString()); | ||
| final GlobalState globalState = mock(GlobalState.class); | ||
| final Map<String, Object> props = Map.of("partitions", partitions); | ||
| when(globalState.getProgressState()).thenReturn(Optional.of(props)); | ||
| when(sourceCoordinator.getPartition(S3PartitionCreatorScheduler.S3_FOLDER_PREFIX + collection)).thenReturn(Optional.of(globalState)); | ||
| Optional<S3PartitionStatus> partitionStatus = s3FolderPartitionCoordinator.getGlobalS3FolderCreationStatus(collection); | ||
| assertThat(partitionStatus.isEmpty(), is(false)); | ||
| assertThat(partitionStatus.get().getPartitions(), is(partitions)); | ||
| } | ||
| } |
75 changes: 75 additions & 0 deletions
75
...org/opensearch/dataprepper/plugins/mongo/s3partition/S3PartitionCreatorSchedulerTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package org.opensearch.dataprepper.plugins.mongo.s3partition; | ||
|
|
||
| import org.junit.jupiter.api.BeforeEach; | ||
| import org.junit.jupiter.api.Test; | ||
| import org.junit.jupiter.api.extension.ExtendWith; | ||
| import org.mockito.ArgumentCaptor; | ||
| import org.mockito.Mock; | ||
| import org.mockito.junit.jupiter.MockitoExtension; | ||
| import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator; | ||
| import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; | ||
| import org.opensearch.dataprepper.plugins.mongo.coordination.partition.GlobalState; | ||
| import org.opensearch.dataprepper.plugins.mongo.coordination.partition.S3FolderPartition; | ||
|
|
||
| import java.time.Duration; | ||
| import java.util.List; | ||
| import java.util.Optional; | ||
| import java.util.Random; | ||
| import java.util.UUID; | ||
| import java.util.concurrent.ExecutorService; | ||
| import java.util.concurrent.Executors; | ||
|
|
||
| import static org.awaitility.Awaitility.await; | ||
| import static org.hamcrest.MatcherAssert.assertThat; | ||
| import static org.hamcrest.Matchers.hasKey; | ||
| import static org.hamcrest.Matchers.hasSize; | ||
| import static org.hamcrest.Matchers.is; | ||
| import static org.mockito.ArgumentMatchers.any; | ||
| import static org.mockito.BDDMockito.given; | ||
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.never; | ||
| import static org.mockito.Mockito.verify; | ||
| import static org.opensearch.dataprepper.plugins.mongo.s3partition.S3PartitionCreatorScheduler.S3_FOLDER_PREFIX; | ||
|
|
||
| @ExtendWith(MockitoExtension.class) | ||
| public class S3PartitionCreatorSchedulerTest { | ||
| @Mock | ||
| private EnhancedSourceCoordinator coordinator; | ||
| private S3PartitionCreatorScheduler s3PartitionCreatorScheduler; | ||
|
|
||
| @BeforeEach | ||
| public void setup() { | ||
| s3PartitionCreatorScheduler = new S3PartitionCreatorScheduler(coordinator, List.of(UUID.randomUUID().toString())); | ||
| } | ||
|
|
||
| @Test | ||
| void test_S3FolderPartition_empty() { | ||
| given(coordinator.acquireAvailablePartition(S3FolderPartition.PARTITION_TYPE)).willReturn(Optional.empty()); | ||
| final ExecutorService executorService = Executors.newSingleThreadExecutor(); | ||
| executorService.submit(() -> s3PartitionCreatorScheduler.run()); | ||
| await() | ||
| .atMost(Duration.ofSeconds(2)) | ||
| .untilAsserted(() -> verify(coordinator, never()).completePartition(any(EnhancedSourcePartition.class))); | ||
| await() | ||
| .atMost(Duration.ofSeconds(2)) | ||
| .untilAsserted(() -> verify(coordinator, never()).createPartition(any(EnhancedSourcePartition.class))); | ||
| executorService.shutdownNow(); | ||
| } | ||
|
|
||
| @Test | ||
| void test_S3FolderPartition_exist() { | ||
| final S3FolderPartition s3FolderPartition = mock(S3FolderPartition.class); | ||
| given(s3FolderPartition.getPartitionCount()).willReturn(Math.abs(new Random().nextInt(100))); | ||
| given(s3FolderPartition.getCollection()).willReturn(UUID.randomUUID().toString()); | ||
| given(coordinator.acquireAvailablePartition(S3FolderPartition.PARTITION_TYPE)).willReturn(Optional.of(s3FolderPartition)); | ||
| s3PartitionCreatorScheduler.run(); | ||
| verify(coordinator).completePartition(s3FolderPartition); | ||
| final ArgumentCaptor<GlobalState> argumentCaptor = ArgumentCaptor.forClass(GlobalState.class); | ||
| verify(coordinator).createPartition(argumentCaptor.capture()); | ||
| final GlobalState globalState = argumentCaptor.getValue(); | ||
| assertThat(globalState.getPartitionKey(), is(S3_FOLDER_PREFIX + s3FolderPartition.getCollection())); | ||
| assertThat(globalState.getProgressState().get(), hasKey("partitions")); | ||
| final List<String> partitions = (List<String>) globalState.getProgressState().get().get("partitions"); | ||
| assertThat(partitions, hasSize(s3FolderPartition.getPartitionCount())); | ||
| } | ||
| } |
20 changes: 20 additions & 0 deletions
20
...est/java/org/opensearch/dataprepper/plugins/mongo/s3partition/S3PartitionCreatorTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| package org.opensearch.dataprepper.plugins.mongo.s3partition; | ||
|
|
||
| import org.junit.jupiter.api.Test; | ||
|
|
||
| import java.util.List; | ||
| import java.util.Random; | ||
|
|
||
| import static org.hamcrest.MatcherAssert.assertThat; | ||
| import static org.hamcrest.Matchers.hasSize; | ||
|
|
||
| public class S3PartitionCreatorTest { | ||
|
|
||
| @Test | ||
| public void createPartitionTest() { | ||
| final int partitionCount = Math.abs(new Random().nextInt(1000)); | ||
| final S3PartitionCreator s3PartitionCreator = new S3PartitionCreator(partitionCount); | ||
| final List<String> partitions = s3PartitionCreator.createPartition(); | ||
| assertThat(partitions, hasSize(partitionCount)); | ||
| } | ||
| } |