code: /*
   * Copyright 2019 Mitsunori Komatsu (komamitsu)
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *      http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.komamitsu.fluency.aws.s3.ingester.sender;
  
  import com.google.common.io.ByteStreams;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.ValueSource;
  import org.komamitsu.fluency.RetryableException;
  import software.amazon.awssdk.auth.credentials.AwsCredentials;
  import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
  import software.amazon.awssdk.core.sync.RequestBody;
  import software.amazon.awssdk.regions.Region;
  import software.amazon.awssdk.services.s3.S3Client;
  import software.amazon.awssdk.services.s3.S3ClientBuilder;
  import software.amazon.awssdk.services.s3.model.PutObjectRequest;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.net.URI;
  import java.nio.ByteBuffer;
  import java.nio.charset.StandardCharsets;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.zip.GZIPInputStream;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.ArgumentMatchers.eq;
  import static org.mockito.Mockito.doAnswer;
  import static org.mockito.Mockito.doReturn;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  
  class AwsS3SenderTest
  {
      @Test
      void buildClientWithDefaults()
      {
          AwsS3Sender.Config config = new AwsS3Sender.Config();
  
          S3Client s3Client = mock(S3Client.class);
          S3ClientBuilder s3ClientBuilder = mock(S3ClientBuilder.class);
          doReturn(s3Client).when(s3ClientBuilder).build();
  
          new AwsS3Sender(s3ClientBuilder, config);
  
          verify(s3ClientBuilder, times(1)).build();
          verify(s3ClientBuilder, times(0)).endpointOverride(any());
          verify(s3ClientBuilder, times(0)).region(any());
          verify(s3ClientBuilder, times(0)).credentialsProvider(any());
      }
  
      @Test
      void buildClientWithCustomizedConfig()
      {
          AwsS3Sender.Config config = new AwsS3Sender.Config();
          config.setEndpoint("https://another.s3endpoi.nt");
          config.setRegion("ap-northeast-1");
          config.setAwsAccessKeyId("foo");
          config.setAwsSecretAccessKey("bar");
  
          S3Client s3Client = mock(S3Client.class);
          S3ClientBuilder s3ClientBuilder = mock(S3ClientBuilder.class);
          doReturn(s3Client).when(s3ClientBuilder).build();
          doAnswer(invocation -> {
              AwsCredentialsProvider provider = invocation.getArgument(0);
              AwsCredentials awsCredentials = provider.resolveCredentials();
              assertEquals("foo", awsCredentials.accessKeyId());
              assertEquals("bar", awsCredentials.secretAccessKey());
              return null;
          }).when(s3ClientBuilder).credentialsProvider(any());
  
          new AwsS3Sender(s3ClientBuilder, config);
  
          verify(s3ClientBuilder, times(1)).build();
          verify(s3ClientBuilder, times(1)).endpointOverride(eq(URI.create("https://another.s3endpoi.nt")));
          verify(s3ClientBuilder, times(1)).region(eq(Region.AP_NORTHEAST_1));
          verify(s3ClientBuilder, times(1)).credentialsProvider(any());
      }
  
      private void testSend(
              AwsS3Sender.Config config,
              boolean gzipCompressed,
              int failures)
              throws IOException
      {
          S3Client s3Client = mock(S3Client.class);
          S3ClientBuilder s3ClientBuilder = mock(S3ClientBuilder.class);
          doReturn(s3Client).when(s3ClientBuilder).build();
  
          AtomicInteger retryCount = new AtomicInteger();
          doAnswer(invocation -> {
              PutObjectRequest request = invocation.getArgument(0);
  
              assertEquals("hello.world", request.bucket());
              assertEquals("2345/01/31/23/59-59-99999.data", request.key());
  
              RequestBody body = invocation.getArgument(1);
              try (InputStream s3In = body.contentStreamProvider().newStream();
                      InputStream in = gzipCompressed ? new GZIPInputStream(s3In) : s3In) {
                  byte[] content = ByteStreams.toByteArray(in);
                  assertEquals("0123456789", new String(content, StandardCharsets.UTF_8));
              }
  
              if (retryCount.getAndIncrement() < failures) {
                  throw new RuntimeException("Something happened");
              }
  
              return null;
          }).when(s3Client).putObject(any(PutObjectRequest.class), any(RequestBody.class));
  
          AwsS3Sender sender = new AwsS3Sender(s3ClientBuilder, config);
  
          sender.send("hello.world", "2345/01/31/23/59-59-99999.data",
                  ByteBuffer.wrap("0123456789".getBytes(StandardCharsets.UTF_8)));
  
          verify(s3Client, times(failures + 1)).putObject(any(PutObjectRequest.class), any(RequestBody.class));
  
          sender.close();
      }
  
      @ParameterizedTest
      @ValueSource(ints={0, 2})
      void send(int failures)
              throws IOException
      {
          AwsS3Sender.Config config = new AwsS3Sender.Config();
          testSend(config, true, failures);
      }
  
      @ParameterizedTest
      @ValueSource(ints={0, 2})
      void sendWithoutCompression(int failures)
              throws IOException
      {
          AwsS3Sender.Config config = new AwsS3Sender.Config();
          config.setCompressionEnabled(false);
  
          testSend(config, false, failures);
      }
  
      @Test
      void sendButRetryOver()
      {
          AwsS3Sender.Config config = new AwsS3Sender.Config();
          config.setRetryMax(2);
          assertThrows(RetryableException.class, () -> testSend(config, true, 3));
      }
  
      @Test
      void close()
      {
          S3Client s3Client = mock(S3Client.class);
          S3ClientBuilder s3ClientBuilder = mock(S3ClientBuilder.class);
          doReturn(s3Client).when(s3ClientBuilder).build();
  
          AwsS3Sender sender = new AwsS3Sender(s3ClientBuilder);
          sender.close();
  
          verify(s3Client, times(1)).close();
      }
  
  }