From 7d038707d28b0533f41e618e58394669b6153e58 Mon Sep 17 00:00:00 2001 From: Chris Helma Date: Mon, 22 Apr 2024 11:50:37 -0500 Subject: [PATCH] Fixed bug in S3Repo handling; added unit tests Signed-off-by: Chris Helma --- RFS/build.gradle | 12 +- RFS/docker/docker-compose.yml | 4 + .../java/com/rfs/ReindexFromSnapshot.java | 2 +- RFS/src/main/java/com/rfs/common/S3Repo.java | 66 +++-- .../test/java/com/rfs/common/S3RepoTest.java | 270 ++++++++++++++++++ 5 files changed, 331 insertions(+), 23 deletions(-) create mode 100644 RFS/src/test/java/com/rfs/common/S3RepoTest.java diff --git a/RFS/build.gradle b/RFS/build.gradle index 17f19ec2d..092ad5897 100644 --- a/RFS/build.gradle +++ b/RFS/build.gradle @@ -38,6 +38,12 @@ dependencies { implementation 'org.apache.lucene:lucene-analyzers-common:8.11.3' implementation 'org.apache.lucene:lucene-backward-codecs:8.11.3' implementation 'software.amazon.awssdk:s3:2.25.16' + + testImplementation 'org.junit.jupiter:junit-jupiter-api:5.10.2' + testImplementation 'org.junit.jupiter:junit-jupiter-params:5.10.2' + testImplementation 'org.mockito:mockito-core:5.11.0' + testImplementation 'org.mockito:mockito-junit-jupiter:5.11.0' + testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.10.2' } application { @@ -110,4 +116,8 @@ task buildDockerImages { } tasks.getByName('composeUp') - .dependsOn(tasks.getByName('buildDockerImages')) \ No newline at end of file + .dependsOn(tasks.getByName('buildDockerImages')) + +test { + useJUnitPlatform() +} \ No newline at end of file diff --git a/RFS/docker/docker-compose.yml b/RFS/docker/docker-compose.yml index aa71ea492..78abc707c 100644 --- a/RFS/docker/docker-compose.yml +++ b/RFS/docker/docker-compose.yml @@ -26,6 +26,10 @@ services: condition: service_started networks: - migrations + environment: + - AWS_ACCESS_KEY_ID=${access_key} + - AWS_SECRET_ACCESS_KEY=${secret_key} + - AWS_SESSION_TOKEN=${session_token} volumes: - snapshotStorage:/snapshots diff --git a/RFS/src/main/java/com/rfs/ReindexFromSnapshot.java b/RFS/src/main/java/com/rfs/ReindexFromSnapshot.java index 2b0a4d733..83b43edfa 100644 --- a/RFS/src/main/java/com/rfs/ReindexFromSnapshot.java +++ b/RFS/src/main/java/com/rfs/ReindexFromSnapshot.java @@ -159,7 +159,7 @@ public static void main(String[] args) throws InterruptedException { if (snapshotDirPath != null) { repo = new FilesystemRepo(snapshotDirPath); } else if (s3RepoUri != null && s3Region != null && s3LocalDirPath != null) { - repo = new S3Repo(s3LocalDirPath, s3RepoUri, s3Region); + repo = S3Repo.create(s3LocalDirPath, s3RepoUri, s3Region); } else if (snapshotLocalRepoDirPath != null) { repo = new FilesystemRepo(snapshotLocalRepoDirPath); } else { diff --git a/RFS/src/main/java/com/rfs/common/S3Repo.java b/RFS/src/main/java/com/rfs/common/S3Repo.java index bd5afc7dc..b502bd92f 100644 --- a/RFS/src/main/java/com/rfs/common/S3Repo.java +++ b/RFS/src/main/java/com/rfs/common/S3Repo.java @@ -11,6 +11,7 @@ import software.amazon.awssdk.core.sync.ResponseTransformer; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.GetObjectRequest; import software.amazon.awssdk.services.s3.model.ListObjectsV2Request; import software.amazon.awssdk.services.s3.model.ListObjectsV2Response; import software.amazon.awssdk.services.s3.model.S3Object; @@ -34,9 +35,9 @@ private static int extractVersion(String key) { } } - private String findRepoFileUri() { - String bucketName = s3RepoUri.split("/")[2]; - String prefix = s3RepoUri.split(bucketName + "/")[1]; + protected String findRepoFileUri() { + String bucketName = getS3BucketName(); + String prefix = getS3ObjectsPrefix(); ListObjectsV2Request listRequest = ListObjectsV2Request.builder() .bucket(bucketName) @@ -46,39 +47,47 @@ private String findRepoFileUri() { ListObjectsV2Response listResponse = s3Client.listObjectsV2(listRequest); Optional highestVersionedIndexFile = listResponse.contents().stream() - .filter(s3Object -> s3Object.key().matches(".*/index-\\d+$")) // Regex to match index files + .filter(s3Object -> s3Object.key().matches(".*index-\\d+$")) // Regex to match index files .max(Comparator.comparingInt(s3Object -> extractVersion(s3Object.key()))); return highestVersionedIndexFile .map(s3Object -> "s3://" + bucketName + "/" + s3Object.key()) - .orElse("No index files found in the specified directory."); + .orElse(null); + } + + protected void ensureS3LocalDirectoryExists(Path localPath) throws IOException { + Files.createDirectories(localPath); } private void downloadFile(String s3Uri, Path localPath) throws IOException { logger.info("Downloading file from S3: " + s3Uri + " to " + localPath); - Files.createDirectories(localPath.getParent()); + ensureS3LocalDirectoryExists(localPath.getParent()); String bucketName = s3Uri.split("/")[2]; String key = s3Uri.split(bucketName + "/")[1]; - s3Client.getObject((req) -> req.bucket(bucketName).key(key), ResponseTransformer.toFile(localPath)); - } + GetObjectRequest getObjectRequest = GetObjectRequest.builder() + .bucket(bucketName) + .key(key) + .build(); - public S3Repo(Path s3LocalDir, String s3Uri, String s3Region) { - this.s3LocalDir = s3LocalDir; + s3Client.getObject(getObjectRequest, ResponseTransformer.toFile(localPath)); + } - // Remove any trailing slash from the S3 URI - if (s3Uri.endsWith("/")) { - this.s3RepoUri = s3Uri.substring(0, s3Uri.length() - 1); - } else { - this.s3RepoUri = s3Uri; - } - - this.s3Region = s3Region; - this.s3Client = S3Client.builder() - .region(Region.of(this.s3Region)) + public static S3Repo create(Path s3LocalDir, String s3Uri, String s3Region) { + S3Client s3Client = S3Client.builder() + .region(Region.of(s3Region)) .credentialsProvider(DefaultCredentialsProvider.create()) .build(); + + return new S3Repo(s3LocalDir, s3Uri, s3Region, s3Client); + } + + public S3Repo(Path s3LocalDir, String s3Uri, String s3Region, S3Client s3Client) { + this.s3LocalDir = s3LocalDir; + this.s3RepoUri = s3Uri; + this.s3Region = s3Region; + this.s3Client = s3Client; } public Path getRepoRootDir() { @@ -88,7 +97,6 @@ public Path getRepoRootDir() { public Path getSnapshotRepoDataFilePath() throws IOException { String repoFileS3Uri = findRepoFileUri(); - // s3://bucket-name/path/to/index-1 => to/index-1 String relativeFileS3Uri = repoFileS3Uri.substring(s3RepoUri.length() + 1); Path localFilePath = s3LocalDir.resolve(relativeFileS3Uri); @@ -137,4 +145,20 @@ public Path getBlobFilePath(String indexId, int shardId, String blobName) throws downloadFile(s3RepoUri + "/" + suffix, filePath); return filePath; } + + public String getS3BucketName() { + String[] parts = s3RepoUri.substring(5).split("/", 2); + return parts[0]; + } + + public String getS3ObjectsPrefix() { + String[] parts = s3RepoUri.substring(5).split("/", 2); + String prefix = parts.length > 1 ? parts[1] : ""; + + if (!prefix.isEmpty() && prefix.endsWith("/")) { + prefix = prefix.substring(0, prefix.length() - 1); + } + + return prefix; + } } diff --git a/RFS/src/test/java/com/rfs/common/S3RepoTest.java b/RFS/src/test/java/com/rfs/common/S3RepoTest.java new file mode 100644 index 000000000..6bef1bfab --- /dev/null +++ b/RFS/src/test/java/com/rfs/common/S3RepoTest.java @@ -0,0 +1,270 @@ +package com.rfs.common; + +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.api.Test; + +import org.mockito.Mockito; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import software.amazon.awssdk.core.sync.ResponseTransformer; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.GetObjectRequest; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.*; + + +@ExtendWith(MockitoExtension.class) +public class S3RepoTest { + @Mock + private S3Client mockS3Client; + private TestableS3Repo testRepo; + private Path testDir = Paths.get("/fake/path"); + private String testRegion = "us-fake-1"; + private String testUri = "s3://bucket-name/directory"; + private String testRepoFileName = "index-2"; + private String testRepoFileUri = testUri + "/" + testRepoFileName; + + + class TestableS3Repo extends S3Repo { + public TestableS3Repo(Path s3LocalDir, String s3RepoUri, String s3Region, S3Client s3Client) { + super(s3LocalDir, s3RepoUri, s3Region, s3Client); + } + + @Override + protected void ensureS3LocalDirectoryExists(Path path) { + // Do nothing + } + + @Override + protected String findRepoFileUri() { + return testRepoFileUri; + } + } + + @BeforeEach + void setUp() { + testRepo = Mockito.spy(new TestableS3Repo(testDir, testUri, testRegion, mockS3Client)); + } + + @Test + void GetRepoRootDir_AsExpected() throws IOException { + // Run the test + Path filePath = testRepo.getRepoRootDir(); + + // Check the results + assertEquals(testDir, filePath); + } + + @Test + void GetSnapshotRepoDataFilePath_AsExpected() throws IOException { + // Set up the test + Path expectedPath = testDir.resolve(testRepoFileName); + String expectedBucketName = testRepo.getS3BucketName(); + String expectedKey = testRepo.getS3ObjectsPrefix() + "/" + testRepoFileName; + + // Run the test + Path filePath = testRepo.getSnapshotRepoDataFilePath(); + + // Check the results + assertEquals(expectedPath, filePath); + + Mockito.verify(testRepo, times(1)).ensureS3LocalDirectoryExists(expectedPath.getParent()); + + GetObjectRequest expectedRequest = GetObjectRequest.builder() + .bucket(expectedBucketName) + .key(expectedKey) + .build(); + + verify(mockS3Client).getObject(eq(expectedRequest), any(ResponseTransformer.class)); + } + + + @Test + void GetGlobalMetadataFilePath_AsExpected() throws IOException { + // Set up the test + String snapshotId = "snapshot1"; + String snapshotFileName = "meta-" + snapshotId + ".dat"; + Path expectedPath = testDir.resolve(snapshotFileName); + String expectedBucketName = testRepo.getS3BucketName(); + String expectedKey = testRepo.getS3ObjectsPrefix() + "/" + snapshotFileName; + + // Run the test + Path filePath = testRepo.getGlobalMetadataFilePath(snapshotId); + + // Check the results + assertEquals(expectedPath, filePath); + + Mockito.verify(testRepo, times(1)).ensureS3LocalDirectoryExists(expectedPath.getParent()); + + GetObjectRequest expectedRequest = GetObjectRequest.builder() + .bucket(expectedBucketName) + .key(expectedKey) + .build(); + + verify(mockS3Client).getObject(eq(expectedRequest), any(ResponseTransformer.class)); + } + + @Test + void GetSnapshotMetadataFilePath_AsExpected() throws IOException { + // Set up the test + String snapshotId = "snapshot1"; + String snapshotFileName = "snap-" + snapshotId + ".dat"; + Path expectedPath = testDir.resolve(snapshotFileName); + String expectedBucketName = testRepo.getS3BucketName(); + String expectedKey = testRepo.getS3ObjectsPrefix() + "/" + snapshotFileName; + + // Run the test + Path filePath = testRepo.getSnapshotMetadataFilePath(snapshotId); + + // Check the results + assertEquals(expectedPath, filePath); + + Mockito.verify(testRepo, times(1)).ensureS3LocalDirectoryExists(expectedPath.getParent()); + + GetObjectRequest expectedRequest = GetObjectRequest.builder() + .bucket(expectedBucketName) + .key(expectedKey) + .build(); + + verify(mockS3Client).getObject(eq(expectedRequest), any(ResponseTransformer.class)); + } + + @Test + void GetIndexMetadataFilePath_AsExpected() throws IOException { + // Set up the test + String indexId = "123abc"; + String indexFileId = "234bcd"; + String indexFileName = "indices/" + indexId + "/meta-" + indexFileId + ".dat"; + Path expectedPath = testDir.resolve(indexFileName); + String expectedBucketName = testRepo.getS3BucketName(); + String expectedKey = testRepo.getS3ObjectsPrefix() + "/" + indexFileName; + + // Run the test + Path filePath = testRepo.getIndexMetadataFilePath(indexId, indexFileId); + + // Check the results + assertEquals(expectedPath, filePath); + + Mockito.verify(testRepo, times(1)).ensureS3LocalDirectoryExists(expectedPath.getParent()); + + GetObjectRequest expectedRequest = GetObjectRequest.builder() + .bucket(expectedBucketName) + .key(expectedKey) + .build(); + + verify(mockS3Client).getObject(eq(expectedRequest), any(ResponseTransformer.class)); + } + + @Test + void GetShardDirPath_AsExpected() throws IOException { + // Set up the test + String indexId = "123abc"; + int shardId = 7; + String shardDirName = "indices/" + indexId + "/" + shardId; + Path expectedPath = testDir.resolve(shardDirName); + + // Run the test + Path filePath = testRepo.getShardDirPath(indexId, shardId); + + // Check the results + assertEquals(expectedPath, filePath); + } + + @Test + void GetShardMetadataFilePath_AsExpected() throws IOException { + // Set up the test + String snapshotId = "snapshot1"; + String indexId = "123abc"; + int shardId = 7; + String shardFileName = "indices/" + indexId + "/" + shardId + "/snap-" + snapshotId + ".dat"; + Path expectedPath = testDir.resolve(shardFileName); + String expectedBucketName = testRepo.getS3BucketName(); + String expectedKey = testRepo.getS3ObjectsPrefix() + "/" + shardFileName; + + // Run the test + Path filePath = testRepo.getShardMetadataFilePath(snapshotId, indexId, shardId); + + // Check the results + assertEquals(expectedPath, filePath); + + Mockito.verify(testRepo, times(1)).ensureS3LocalDirectoryExists(expectedPath.getParent()); + + GetObjectRequest expectedRequest = GetObjectRequest.builder() + .bucket(expectedBucketName) + .key(expectedKey) + .build(); + + verify(mockS3Client).getObject(eq(expectedRequest), any(ResponseTransformer.class)); + } + + @Test + void GetBlobFilePath_AsExpected() throws IOException { + // Set up the test + String blobName = "bobloblaw"; + String indexId = "123abc"; + int shardId = 7; + String shardFileName = "indices/" + indexId + "/" + shardId + "/" + blobName; + Path expectedPath = testDir.resolve(shardFileName); + String expectedBucketName = testRepo.getS3BucketName(); + String expectedKey = testRepo.getS3ObjectsPrefix() + "/" + shardFileName; + + // Run the test + Path filePath = testRepo.getBlobFilePath(indexId, shardId, blobName); + + // Check the results + assertEquals(expectedPath, filePath); + + Mockito.verify(testRepo, times(1)).ensureS3LocalDirectoryExists(expectedPath.getParent()); + + GetObjectRequest expectedRequest = GetObjectRequest.builder() + .bucket(expectedBucketName) + .key(expectedKey) + .build(); + + verify(mockS3Client).getObject(eq(expectedRequest), any(ResponseTransformer.class)); + } + + static Stream provideUrisForBucketNames() { + return Stream.of( + Arguments.of("s3://bucket-name", "bucket-name"), + Arguments.of("s3://bucket-name/", "bucket-name"), + Arguments.of("s3://bucket-name/with/suffix", "bucket-name") + ); + } + + @ParameterizedTest + @MethodSource("provideUrisForBucketNames") + void getS3BucketName_AsExpected(String uri, String expectedBucketName) { + TestableS3Repo repo = new TestableS3Repo(testDir, uri, testRegion, mock(S3Client.class)); + assertEquals(expectedBucketName, repo.getS3BucketName()); + } + + static Stream provideUrisForPrefixes() { + return Stream.of( + Arguments.of("s3://bucket-name", ""), + Arguments.of("s3://bucket-name/", ""), + Arguments.of("s3://bucket-name/with/suffix", "with/suffix"), + Arguments.of("s3://bucket-name/with/suffix/", "with/suffix") + ); + } + + @ParameterizedTest + @MethodSource("provideUrisForPrefixes") + void getS3ObjectsPrefix_AsExpected(String uri, String expectedPrefix) { + TestableS3Repo repo = new TestableS3Repo(testDir, uri, testRegion, mock(S3Client.class)); + assertEquals(expectedPrefix, repo.getS3ObjectsPrefix()); + } +}