Add more unit tests

extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/S3TaskLogsTest.java

@@ -35,8 +35,11 @@
 import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import org.apache.commons.io.IOUtils;
 import org.apache.druid.common.utils.CurrentTimeMillisSupplier;
 import org.apache.druid.java.util.common.StringUtils;
+import org.easymock.Capture;
+import org.easymock.CaptureType;
 import org.easymock.EasyMock;
 import org.easymock.EasyMockRunner;
 import org.easymock.EasyMockSupport;
@@ -55,6 +58,7 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.net.URI;
+import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -143,6 +147,78 @@ public void test_pushTaskStatus() throws IOException
     EasyMock.verify(s3Client);
   }
 
+  @Test
+  public void test_pushTaskPayload() throws IOException
+  {
+    Capture<PutObjectRequest> putObjectRequestCapture = Capture.newInstance(CaptureType.FIRST);
+    EasyMock.expect(s3Client.putObject(EasyMock.capture(putObjectRequestCapture)))
+        .andReturn(new PutObjectResult())
+        .once();
+
+    EasyMock.replay(s3Client);
+
+    S3TaskLogsConfig config = new S3TaskLogsConfig();
+    config.setS3Bucket(TEST_BUCKET);
+    config.setS3Prefix("prefix");
+    config.setDisableAcl(true);
+
+    CurrentTimeMillisSupplier timeSupplier = new CurrentTimeMillisSupplier();
+    S3InputDataConfig inputDataConfig = new S3InputDataConfig();
+    S3TaskLogs s3TaskLogs = new S3TaskLogs(s3Client, config, inputDataConfig, timeSupplier);
+
+    File payloadFile = tempFolder.newFile("task.json");
+    String taskId = "index_test-datasource_2019-06-18T13:30:28.887Z";
+    s3TaskLogs.pushTaskPayload(taskId, payloadFile);
+
+    PutObjectRequest putObjectRequest = putObjectRequestCapture.getValue();
+    Assert.assertEquals(TEST_BUCKET, putObjectRequest.getBucketName());
+    Assert.assertEquals("prefix/" + taskId + "/task.json", putObjectRequest.getKey());
+    Assert.assertEquals(payloadFile, putObjectRequest.getFile());
+    EasyMock.verify(s3Client);
+  }
+
+  @Test
+  public void test_streamTaskPayload() throws IOException
+  {
+    String taskPayloadString = "task payload";
+
+    ObjectMetadata objectMetadata = new ObjectMetadata();
+    objectMetadata.setContentLength(taskPayloadString.length());
+    EasyMock.expect(s3Client.getObjectMetadata(EasyMock.anyObject(), EasyMock.anyObject()))
+        .andReturn(objectMetadata)
+        .once();
+
+    InputStream taskPayload = new ByteArrayInputStream(taskPayloadString.getBytes(Charset.defaultCharset()));
+    S3Object s3Object = new S3Object();
+    s3Object.setObjectContent(taskPayload);
+    Capture<GetObjectRequest> getObjectRequestCapture = Capture.newInstance(CaptureType.FIRST);
+    EasyMock.expect(s3Client.getObject(EasyMock.capture(getObjectRequestCapture)))
+        .andReturn(s3Object)
+        .once();
+
+    EasyMock.replay(s3Client);
+
+    S3TaskLogsConfig config = new S3TaskLogsConfig();
+    config.setS3Bucket(TEST_BUCKET);
+    config.setS3Prefix("prefix");
+    config.setDisableAcl(true);
+
+    CurrentTimeMillisSupplier timeSupplier = new CurrentTimeMillisSupplier();
+    S3InputDataConfig inputDataConfig = new S3InputDataConfig();
+    S3TaskLogs s3TaskLogs = new S3TaskLogs(s3Client, config, inputDataConfig, timeSupplier);
+
+    String taskId = "index_test-datasource_2019-06-18T13:30:28.887Z";
+    Optional<InputStream> payloadResponse = s3TaskLogs.streamTaskPayload(taskId);
+
+    GetObjectRequest getObjectRequest = getObjectRequestCapture.getValue();
+    Assert.assertEquals(TEST_BUCKET, getObjectRequest.getBucketName());
+    Assert.assertEquals("prefix/" + taskId + "/task.json", getObjectRequest.getKey());
+    Assert.assertTrue(payloadResponse.isPresent());
+
+    Assert.assertEquals(taskPayloadString, IOUtils.toString(payloadResponse.get(), Charset.defaultCharset()));
+    EasyMock.verify(s3Client);
+  }
+
   @Test
   public void test_killAll_noException_deletesAllTaskLogs() throws IOException
   {