cameronlee314 commented on a change in pull request #1239: SAMZA-2421: Add SystemProducer for Azure Blob Storage URL: https://github.com/apache/samza/pull/1239#discussion_r360997259
########## File path: samza-azure/src/test/java/org/apache/samza/system/azureblob/avro/TestAzureBlobOutputStream.java ########## @@ -0,0 +1,347 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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.apache.samza.system.azureblob.avro; + +import com.azure.core.http.rest.SimpleResponse; +import com.azure.core.implementation.util.FluxUtil; +import org.apache.samza.system.azureblob.compression.Compression; +import org.apache.samza.system.azureblob.producer.AzureBlobWriterMetrics; +import com.azure.storage.blob.specialized.BlockBlobAsyncClient; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Map; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.ArgumentCaptor; +import org.powermock.api.mockito.PowerMockito; +import org.powermock.core.classloader.annotations.PrepareForTest; +import org.powermock.modules.junit4.PowerMockRunner; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.anyList; +import static org.mockito.Mockito.anyLong; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@RunWith(PowerMockRunner.class) +@PrepareForTest({BlockBlobAsyncClient.class}) +public class TestAzureBlobOutputStream { + private ThreadPoolExecutor threadPool; + private ByteArrayOutputStream mockByteArrayOutputStream; + private static final int THRESHOLD = 100; + private BlockBlobAsyncClient mockBlobAsyncClient; + private AzureBlobOutputStream azureBlobOutputStream; + private static String randomString = "roZzozzLiR7GCEjcB0UsRUNgBAip8cSLGXQSo3RQvbIDoxOaaRs4hrec2s5rMPWgTPRY4UnE959worEtyhRjwUFnRnVuNFZ554yuPQCbI69qFkQX7MmrB4blmpSnFeGjWKjFjIRLFNVSsQBYMkr5jT4T83uVtuGumsjACVrpcilihdd194H8Y71rQcrXZoTQtw5OvmPicbwptawpHoRNzHihyaDVYgAs0dQbvVEu1gitKpamzYdMLFtc5h8PFZSVEB"; + private static byte[] bytes = randomString.substring(0, THRESHOLD).getBytes(); + private static byte[] compressedBytes = randomString.substring(0, THRESHOLD / 2).getBytes(); + private AzureBlobWriterMetrics mockMetrics; + private Compression mockCompression; + + @Before + public void setup() throws Exception { + threadPool = new ThreadPoolExecutor(1, 1, 60, + TimeUnit.SECONDS, new LinkedBlockingDeque<Runnable>()); + + + mockByteArrayOutputStream = spy(new ByteArrayOutputStream(THRESHOLD)); + + mockBlobAsyncClient = PowerMockito.mock(BlockBlobAsyncClient.class); + when(mockBlobAsyncClient.stageBlock(anyString(), any(), anyLong())).thenReturn( + Mono.just(new SimpleResponse(null, 200, null, null)).flatMap(FluxUtil::toMono)); + when(mockBlobAsyncClient.commitBlockListWithResponse(any(), any(), any(), any(), any())).thenReturn( + Mono.just(new SimpleResponse(null, 200, null, null))); + + when(mockBlobAsyncClient.getBlobUrl()).thenReturn("https://samza.blob.core.windows.net/fake-blob-url";); + + mockMetrics = mock(AzureBlobWriterMetrics.class); + + mockCompression = mock(Compression.class); + doReturn(compressedBytes).when(mockCompression).compress(bytes); + + azureBlobOutputStream = spy(new AzureBlobOutputStream(mockBlobAsyncClient, threadPool, mockMetrics, + 60000, THRESHOLD, mockByteArrayOutputStream, mockCompression)); + } + + @Test + public void testWrite() { + byte[] b = new byte[THRESHOLD - 10]; + azureBlobOutputStream.write(b, 0, THRESHOLD - 10); + verify(mockBlobAsyncClient, never()).stageBlock(any(), any(), anyLong()); // since size of byte[] written is less than threshold + verify(mockMetrics).updateWriteByteMetrics(THRESHOLD - 10); + verify(mockMetrics, never()).updateAzureUploadMetrics(); + } + + @Test + public void testWriteLargerThanThreshold() { + byte[] largeRecord = randomString.substring(0, 2 * THRESHOLD).getBytes(); + byte[] largeRecordFirstHalf = randomString.substring(0, THRESHOLD).getBytes(); + byte[] largeRecordSecondHalf = randomString.substring(THRESHOLD, 2 * THRESHOLD).getBytes(); + + byte[] compressB = randomString.substring(0, THRESHOLD / 2).getBytes(); + + doReturn(compressB).when(mockCompression).compress(largeRecordFirstHalf); + doReturn(compressB).when(mockCompression).compress(largeRecordSecondHalf); + + azureBlobOutputStream.write(largeRecord, 0, 2 * THRESHOLD); + // azureBlobOutputStream.close waits on the CompletableFuture which does the actual stageBlock in uploadBlockAsync + azureBlobOutputStream.close(); + + // invoked 2 times for the data which is 2*threshold + verify(mockCompression).compress(largeRecordFirstHalf); + verify(mockCompression).compress(largeRecordSecondHalf); + ArgumentCaptor<Flux> argument = ArgumentCaptor.forClass(Flux.class); + verify(mockBlobAsyncClient).stageBlock(eq(blockIdEncoded(0)), argument.capture(), eq((long) compressB.length)); + verify(mockBlobAsyncClient).stageBlock(eq(blockIdEncoded(1)), argument.capture(), eq((long) compressB.length)); + argument.getAllValues().forEach(flux -> { + Assert.assertEquals(ByteBuffer.wrap(compressB), flux.blockFirst()); + }); + verify(mockMetrics).updateWriteByteMetrics(2 * THRESHOLD); + verify(mockMetrics, times(2)).updateAzureUploadMetrics(); + } + + @Test + public void testWriteLargeRecordWithSmallRecordInBuffer() { + byte[] halfBlock = new byte[THRESHOLD / 2]; + byte[] fullBlock = new byte[THRESHOLD]; + byte[] largeRecord = new byte[2 * THRESHOLD]; + byte[] fullBlockCompressedByte = new byte[50]; + byte[] halfBlockCompressedByte = new byte[25]; + doReturn(fullBlockCompressedByte).when(mockCompression).compress(fullBlock); + doReturn(halfBlockCompressedByte).when(mockCompression).compress(halfBlock); + + // FIRST write a small record = same as half block + azureBlobOutputStream.write(halfBlock, 0, THRESHOLD / 2); + verify(mockMetrics).updateWriteByteMetrics(THRESHOLD / 2); + + // SECOND write the large record + azureBlobOutputStream.write(largeRecord, 0, 2 * THRESHOLD); + verify(mockMetrics).updateWriteByteMetrics(2 * THRESHOLD); + + azureBlobOutputStream.flush(); // to flush out buffered data + + // azureBlobOutputStream.close waits on the CompletableFuture which does the actual stageBlock in uploadBlockAsync + azureBlobOutputStream.close(); + + verify(mockCompression, times(2)).compress(fullBlock); + verify(mockCompression).compress(halfBlock); + + ArgumentCaptor<Flux> argument = ArgumentCaptor.forClass(Flux.class); + ArgumentCaptor<Flux> argument2 = ArgumentCaptor.forClass(Flux.class); + verify(mockBlobAsyncClient).stageBlock(eq(blockIdEncoded(0)), argument.capture(), eq((long) fullBlockCompressedByte.length)); + verify(mockBlobAsyncClient).stageBlock(eq(blockIdEncoded(1)), argument.capture(), eq((long) fullBlockCompressedByte.length)); + verify(mockBlobAsyncClient).stageBlock(eq(blockIdEncoded(2)), argument2.capture(), eq((long) halfBlockCompressedByte.length)); + argument.getAllValues().forEach(flux -> { + Assert.assertEquals(ByteBuffer.wrap(fullBlockCompressedByte), flux.blockFirst()); + }); + Assert.assertEquals(ByteBuffer.wrap(halfBlockCompressedByte), ((Flux) argument2.getValue()).blockFirst()); + verify(mockMetrics, times(3)).updateAzureUploadMetrics(); + } + + + @Test + public void testWriteThresholdCrossed() throws Exception { + azureBlobOutputStream.write(bytes, 0, THRESHOLD / 2); + azureBlobOutputStream.write(bytes, THRESHOLD / 2, THRESHOLD / 2); + // azureBlobOutputStream.close waits on the CompletableFuture which does the actual stageBlock in uploadBlockAsync + azureBlobOutputStream.close(); + + verify(mockCompression).compress(bytes); + ArgumentCaptor<Flux> argument = ArgumentCaptor.forClass(Flux.class); + verify(mockBlobAsyncClient).stageBlock(eq(blockIdEncoded(0)), argument.capture(), eq((long) compressedBytes.length)); // since size of byte[] written is less than threshold + Assert.assertEquals(ByteBuffer.wrap(compressedBytes), ((Flux) argument.getValue()).blockFirst()); + verify(mockMetrics, times(2)).updateWriteByteMetrics(THRESHOLD / 2); + verify(mockMetrics, times(1)).updateAzureUploadMetrics(); + } + + @Test(expected = RuntimeException.class) + public void testWriteFailed() { + when(mockBlobAsyncClient.stageBlock(anyString(), any(), anyLong())).thenThrow(new Exception("Test Failed")); + + byte[] b = new byte[100]; + azureBlobOutputStream.write(b, 0, THRESHOLD); // threshold crossed so stageBlock is scheduled. + // azureBlobOutputStream.close waits on the CompletableFuture which does the actual stageBlock in uploadBlockAsync + azureBlobOutputStream.close(); + } + + @Test + public void testClose() { + azureBlobOutputStream.write(bytes, 0, THRESHOLD); + int blockNum = 0; + String blockId = String.format("%05d", blockNum); + String blockIdEncoded = Base64.getEncoder().encodeToString(blockId.getBytes()); + + PowerMockito.doAnswer(invocation -> { + ArrayList<String> blockListArg = (ArrayList<String>) invocation.getArguments()[0]; + String blockIdArg = (String) blockListArg.toArray()[0]; + Assert.assertEquals(blockIdEncoded, blockIdArg); + Map<String, String> blobMetadata = (Map<String, String>) invocation.getArguments()[2]; + Assert.assertEquals(blobMetadata.get(AzureBlobOutputStream.BLOB_RAW_SIZE_BYTES_METADATA), Long.toString(THRESHOLD)); + return Mono.just(new SimpleResponse(null, 200, null, null)); + }).when(mockBlobAsyncClient).commitBlockListWithResponse(anyList(), any(), any(), any(), any()); + + azureBlobOutputStream.close(); + verify(mockMetrics).updateAzureCommitMetrics(); + } + + @Test + public void testCloseMultipleBlocks() { + azureBlobOutputStream.write(bytes, 0, THRESHOLD); + azureBlobOutputStream.write(bytes, 0, THRESHOLD); + + int blockNum = 0; + String blockId = String.format("%05d", blockNum); + String blockIdEncoded = Base64.getEncoder().encodeToString(blockId.getBytes()); + + int blockNum1 = 1; + String blockId1 = String.format("%05d", blockNum1); + String blockIdEncoded1 = Base64.getEncoder().encodeToString(blockId1.getBytes()); + + PowerMockito.doAnswer(invocation -> { + ArrayList<String> blockListArg = (ArrayList<String>) invocation.getArguments()[0]; + String blockIdArg = (String) blockListArg.toArray()[0]; + String blockIdArg1 = (String) blockListArg.toArray()[1]; + Assert.assertEquals(blockIdEncoded, blockIdArg); + Assert.assertEquals(blockIdEncoded1, blockIdArg1); Review comment: Can this be replaced with `Assert.assertEquals(Arrays.asList(blockIdEncoded, blockIdEncoded1), blockListArg)`? ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services
