From 2172aeed0eeb496107115737e5d3060c8d2289ab Mon Sep 17 00:00:00 2001 From: Chris Helma <25470211+chelma@users.noreply.github.com> Date: Tue, 23 Apr 2024 05:42:15 -0500 Subject: [PATCH] Updated RFS Docker Compose to handle AWS Creds; fix S3Repo bug (#597) * Updated RFS README for using AWS Creds in Docker Compose Signed-off-by: Chris Helma * Fixed bug in S3Repo handling; added unit tests Signed-off-by: Chris Helma --------- Signed-off-by: Chris Helma --- RFS/README.md | 91 +++--- RFS/build.gradle | 12 +- RFS/docker/docker-compose.yml | 7 + .../java/com/rfs/ReindexFromSnapshot.java | 2 +- RFS/src/main/java/com/rfs/common/S3Repo.java | 66 +++-- .../test/java/com/rfs/common/S3RepoTest.java | 270 ++++++++++++++++++ 6 files changed, 369 insertions(+), 79 deletions(-) create mode 100644 RFS/src/test/java/com/rfs/common/S3RepoTest.java diff --git a/RFS/README.md b/RFS/README.md index 70bc44c04..b731048a7 100644 --- a/RFS/README.md +++ b/RFS/README.md @@ -7,10 +7,10 @@ - [Using an existing S3 snapshot](#using-an-existing-s3-snapshot) - [Using a source cluster](#using-a-source-cluster) - [Using Docker](#using-docker) + - [Providing AWS permissions for S3 snapshot creation](#providing-aws-permissions-for-s3-snapshot-creation) - [Handling auth](#handling-auth) - [How to set up an ES 6.8 Source Cluster w/ an attached debugger](#how-to-set-up-an-es-68-source-cluster-w-an-attached-debugger) - [How to set up an ES 7.10 Source Cluster running in Docker](#how-to-set-up-an-es-710-source-cluster-running-in-docker) - - [Providing AWS permissions for S3 snapshot creation](#providing-aws-permissions-for-s3-snapshot-creation) - [Setting up the Cluster w/ some sample docs](#setting-up-the-cluster-w-some-sample-docs) - [How to set up an OS 2.11 Target Cluster](#how-to-set-up-an-os-211-target-cluster) @@ -81,8 +81,8 @@ Also built into this Docker/Gradle support is the ability to spin up a testing R This environment can be spun up with the Gradle command, and use the optional `-Pdataset` flag to preload a dataset from the `generateDatasetStage` in the multi-stage Docker [here](docker/TestSource_ES_7_10/Dockerfile). This stage will take a few minutes to run on its first attempt if it is generating data, as it will be making requests with OSB. This will be cached for future runs. ```shell ./gradlew composeUp -Pdataset=default_osb_test_workloads - ``` + And deleted with the Gradle command ```shell ./gradlew composeDown @@ -91,10 +91,10 @@ And deleted with the Gradle command After the Docker compose containers are created the elasticsearch/opensearch source and target clusters can be interacted with like normal. For RFS testing, a user can also load custom templates/indices/documents into the source cluster before kicking off RFS. ```shell # To check indices on the source cluster -curl http://localhost:19200/_cat/indices?v +curl 'http://localhost:19200/_cat/indices?v' # To check indices on the target cluster -curl http://localhost:29200/_cat/indices?v +curl 'http://localhost:29200/_cat/indices?v' ``` To kick off RFS: @@ -102,6 +102,36 @@ To kick off RFS: docker exec -it rfs-compose-reindex-from-snapshot-1 sh -c "/rfs-app/runJavaWithClasspath.sh com.rfs.ReindexFromSnapshot --snapshot-name test-snapshot --snapshot-local-repo-dir /snapshots --min-replicas 0 --lucene-dir '/lucene' --source-host http://elasticsearchsource:9200 --target-host http://opensearchtarget:9200 --source-version es_7_10 --target-version os_2_11" ``` +#### Providing AWS permissions for S3 snapshot creation + +While the source cluster's container will have the `repository-s3` plugin installed out-of-the-box, to use it you'll need to provide AWS Credentials. This plugin will either accept credential [from the Elasticsearch Keystore](https://www.elastic.co/guide/en/elasticsearch/plugins/7.10/repository-s3-client.html) or via the standard ENV variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN`). The issue w/ either approach in local testing is renewal of the timeboxed creds. One possible solution is to use an IAM User, but that is generally frowned upon. The approach we'll take here is to accept that the test cluster is temporary, so the creds can be as well. Therefore, we can make an AWS IAM Role in our AWS Account with the creds it needs, assume it locally to generate the credential triple, and pipe that into the container using ENV variables. + +Start by making an AWS IAM Role (e.g. `arn:aws:iam::XXXXXXXXXXXX:role/testing-es-source-cluster-s3-access`) with S3 Full Access permissions in your AWS Account. You can then get credentials with that identity good for up to one hour: + +```shell +unset access_key && unset secret_key && unset session_token + +output=$(aws sts assume-role --role-arn "arn:aws:iam::XXXXXXXXXXXX:role/testing-es-source-cluster-s3-access" --role-session-name "ES-Source-Cluster") + +export access_key=$(echo $output | jq -r .Credentials.AccessKeyId) +export secret_key=$(echo $output | jq -r .Credentials.SecretAccessKey) +export session_token=$(echo $output | jq -r .Credentials.SessionToken) +``` + +The one hour limit is annoying but workable, given the only thing it's needed for is creating the snapshot at the very start of the RFS process. This is primarily driven by the fact that IAM limits session durations to one hour when the role is assumed via another role (e.g. role chaining). If your original creds in the AWS keyring are from an IAM User, etc, then this might not be a restriction for you and you can have up to 12 hours with the assumed creds. Ideas on how to improve this would be greatly appreciated. + +Anyways, we pipe those ENV variables into the source cluster's container via the Docker Compose file, so you can just launch the test setup as normal: + +```shell +./gradlew composeUp -Pdataset=default_osb_test_workloads +``` + +If you need to renew the creds, you can just kill the existing source container, renew the creds, and spin up a new container. + +``` +./gradlew composeDown +``` + ### Handling auth RFS currently supports both basic auth (username/password) and no auth for both the source and target clusters. To use the no-auth approach, just neglect the username/password arguments. @@ -201,58 +231,7 @@ curl -X PUT "localhost:9200/_snapshot/fs_repository/global_state_snapshot?wait_f ## How to set up an ES 7.10 Source Cluster running in Docker -The `./docker` directory contains a Dockerfile for an ES 7.10 cluster, which you can use for testing. You can run it like so: - -``` -(cd ./docker/TestSource_ES_7_10; docker build . -t es-w-s3) - -docker run -d \ --p 9200:9200 \ --e discovery.type=single-node \ ---name elastic-source \ -es-w-s3 \ -/bin/sh -c '/usr/local/bin/docker-entrypoint.sh eswrapper & wait -n' - -curl http://localhost:9200 -``` - -### Providing AWS permissions for S3 snapshot creation - -While the container will have the `repository-s3` plugin installed out-of-the-box, to use it you'll need to provide AWS Credentials. This plugin will either accept credential [from the Elasticsearch Keystore](https://www.elastic.co/guide/en/elasticsearch/plugins/7.10/repository-s3-client.html) or via the standard ENV variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN`). The issue w/ either approach in local testing is renewal of the timeboxed creds. One possible solution is to use an IAM User, but that is generally frowned upon. The approach we'll take here is to accept that the test cluster is temporary, so the creds can be as well. Therefore, we can make an AWS IAM Role in our AWS Account with the creds it needs, assume it locally to generate the credential triple, and pipe that into the container using ENV variables. - -Start by making an AWS IAM Role (e.g. `arn:aws:iam::XXXXXXXXXXXX:role/testing-es-source-cluster-s3-access`) with S3 Full Access permissions in your AWS Account. You can then get credentials with that identity good for up to one hour: - -``` -unset access_key && unset secret_key && unset session_token - -output=$(aws sts assume-role --role-arn "arn:aws:iam::XXXXXXXXXXXX:role/testing-es-source-cluster-s3-access" --role-session-name "ES-Source-Cluster") - -access_key=$(echo $output | jq -r .Credentials.AccessKeyId) -secret_key=$(echo $output | jq -r .Credentials.SecretAccessKey) -session_token=$(echo $output | jq -r .Credentials.SessionToken) -``` - -The one hour limit is annoying but workable, given the only thing it's needed for is creating the snapshot at the very start of the RFS process. This is primarily driven by the fact that IAM limits session durations to one hour when the role is assumed via another role (e.g. role chaining). If your original creds in the AWS keyring are from an IAM User, etc, then this might not be a restriction for you and you can have up to 12 hours with the assumed creds. Ideas on how to improve this would be greatly appreciated. - -Anyways, you can then launch the container with those temporary credentials like so, using the : - -``` -docker run -d \ --p 9200:9200 \ --e discovery.type=single-node \ --e AWS_ACCESS_KEY_ID=$access_key \ --e AWS_SECRET_ACCESS_KEY=$secret_key \ --e AWS_SESSION_TOKEN=$session_token \ --v ~/.aws:/root/.aws:ro \ ---name elastic-source \ -es-w-s3 -``` - -If you need to renew the creds, you can just kill the existing source container, renew the creds, and spin up a new container. - -``` -docker stop elastic-source; docker rm elastic-source -``` +It is recommended to use the Docker Compose setup described above in #using-docker ### Setting up the Cluster w/ some sample docs 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 74724af29..78abc707c 100644 --- a/RFS/docker/docker-compose.yml +++ b/RFS/docker/docker-compose.yml @@ -8,6 +8,9 @@ services: environment: - discovery.type=single-node - path.repo=/snapshots + - AWS_ACCESS_KEY_ID=${access_key} + - AWS_SECRET_ACCESS_KEY=${secret_key} + - AWS_SESSION_TOKEN=${session_token} ports: - '19200:9200' volumes: @@ -23,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 5a7b45562..1cd897ffc 100644 --- a/RFS/src/main/java/com/rfs/ReindexFromSnapshot.java +++ b/RFS/src/main/java/com/rfs/ReindexFromSnapshot.java @@ -163,7 +163,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()); + } +}