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_r360996977
########## 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]; Review comment: Would it be good to fill in these blocks with some data to be able to verify that the expected data gets written? The test currently just checks the size of the blocks through the byte array comparisons, but doesn't verify that the expected data is written as expected. ---------------------------------------------------------------- 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
