use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.
the class TestAzureBlobAvroWriter method testRecordLimitExceeded.
@Test
public void testRecordLimitExceeded() throws Exception {
String blobUrlPrefix = "test";
String blobNameRegex = "test/[0-9]{4}/[0-9]{2}/[0-9]{2}/[0-9]{2}/[0-9]{2}-[0-9]{2}-.{8}.avro.gz";
AzureBlobWriterMetrics mockMetrics = mock(AzureBlobWriterMetrics.class);
long maxBlobSize = AzureBlobAvroWriter.DATAFILEWRITER_OVERHEAD + 1000;
long maxRecordsPerBlob = 10;
BlobContainerAsyncClient mockContainerClient = PowerMockito.mock(BlobContainerAsyncClient.class);
azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerClient, mockMetrics, threadPool, THRESHOLD, 60000, blobUrlPrefix, null, null, null, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, maxBlobSize, maxRecordsPerBlob, mockCompression, true));
DataFileWriter mockDataFileWriter1 = mock(DataFileWriter.class);
PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter1);
BlobAsyncClient mockBlobAsyncClient1 = mock(BlobAsyncClient.class);
doReturn(mockBlobAsyncClient1).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
BlockBlobAsyncClient mockBlockBlobAsyncClient1 = mock(BlockBlobAsyncClient.class);
doReturn(mockBlockBlobAsyncClient1).when(mockBlobAsyncClient1).getBlockBlobAsyncClient();
AzureBlobOutputStream mockAzureBlobOutputStream1 = mock(AzureBlobOutputStream.class);
PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient1, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream1);
when(mockAzureBlobOutputStream1.getSize()).thenReturn((long) 1);
for (int i = 0; i < maxRecordsPerBlob; i++) {
azureBlobAvroWriter.write(ome);
}
OutgoingMessageEnvelope ome2 = createOME("Topic2");
DataFileWriter mockDataFileWriter2 = mock(DataFileWriter.class);
PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter2);
BlobAsyncClient mockBlobAsyncClient2 = mock(BlobAsyncClient.class);
doReturn(mockBlobAsyncClient2).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
BlockBlobAsyncClient mockBlockBlobAsyncClient2 = mock(BlockBlobAsyncClient.class);
doReturn(mockBlockBlobAsyncClient2).when(mockBlobAsyncClient2).getBlockBlobAsyncClient();
AzureBlobOutputStream mockAzureBlobOutputStream2 = mock(AzureBlobOutputStream.class);
PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient2, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream2);
when(mockAzureBlobOutputStream2.getSize()).thenReturn((long) 1);
azureBlobAvroWriter.write(ome2);
ArgumentCaptor<String> argument = ArgumentCaptor.forClass(String.class);
verify(mockContainerClient, times(2)).getBlobAsyncClient(argument.capture());
argument.getAllValues().forEach(blobName -> {
Assert.assertTrue(blobName.contains(blobUrlPrefix));
});
List<String> allBlobNames = argument.getAllValues();
Assert.assertNotEquals(allBlobNames.get(0), allBlobNames.get(1));
verify(mockDataFileWriter1, times((int) maxRecordsPerBlob)).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome.getMessage())));
verify(mockDataFileWriter2).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome2.getMessage())));
verify(mockDataFileWriter1).create(((IndexedRecord) ome.getMessage()).getSchema(), mockAzureBlobOutputStream1);
verify(mockDataFileWriter2).create(((IndexedRecord) ome2.getMessage()).getSchema(), mockAzureBlobOutputStream2);
}
use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.
the class TestAzureBlobAvroWriter method testMaxBlobSizeExceeded.
@Test
public void testMaxBlobSizeExceeded() throws Exception {
String blobUrlPrefix = "test";
String blobNameRegex = "test/[0-9]{4}/[0-9]{2}/[0-9]{2}/[0-9]{2}/[0-9]{2}-[0-9]{2}-.{8}.avro.gz";
long maxBlobSize = 1000;
AzureBlobWriterMetrics mockMetrics = mock(AzureBlobWriterMetrics.class);
BlobContainerAsyncClient mockContainerClient = PowerMockito.mock(BlobContainerAsyncClient.class);
azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerClient, mockMetrics, threadPool, THRESHOLD, 60000, blobUrlPrefix, null, null, null, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, maxBlobSize, 10, mockCompression, true));
DataFileWriter mockDataFileWriter1 = mock(DataFileWriter.class);
PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter1);
BlobAsyncClient mockBlobAsyncClient1 = mock(BlobAsyncClient.class);
doReturn(mockBlobAsyncClient1).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
BlockBlobAsyncClient mockBlockBlobAsyncClient1 = mock(BlockBlobAsyncClient.class);
doReturn(mockBlockBlobAsyncClient1).when(mockBlobAsyncClient1).getBlockBlobAsyncClient();
AzureBlobOutputStream mockAzureBlobOutputStream1 = mock(AzureBlobOutputStream.class);
PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient1, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream1);
when(mockAzureBlobOutputStream1.getSize()).thenReturn((long) maxBlobSize - 1);
// first OME creates the first blob
azureBlobAvroWriter.write(ome);
OutgoingMessageEnvelope ome2 = createOME("Topic2");
DataFileWriter mockDataFileWriter2 = mock(DataFileWriter.class);
PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter2);
BlobAsyncClient mockBlobAsyncClient2 = mock(BlobAsyncClient.class);
doReturn(mockBlobAsyncClient2).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
BlockBlobAsyncClient mockBlockBlobAsyncClient2 = mock(BlockBlobAsyncClient.class);
doReturn(mockBlockBlobAsyncClient2).when(mockBlobAsyncClient2).getBlockBlobAsyncClient();
AzureBlobOutputStream mockAzureBlobOutputStream2 = mock(AzureBlobOutputStream.class);
PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient2, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream2);
when(mockAzureBlobOutputStream2.getSize()).thenReturn((long) maxBlobSize - 1);
// Second OME creates the second blob because maxBlobSize is 1000 and mockAzureBlobOutputStream.getSize is 999.
azureBlobAvroWriter.write(ome2);
ArgumentCaptor<String> argument = ArgumentCaptor.forClass(String.class);
verify(mockContainerClient, times(2)).getBlobAsyncClient(argument.capture());
argument.getAllValues().forEach(blobName -> {
Assert.assertTrue(blobName.contains(blobUrlPrefix));
});
List<String> allBlobNames = argument.getAllValues();
Assert.assertNotEquals(allBlobNames.get(0), allBlobNames.get(1));
verify(mockDataFileWriter1).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome.getMessage())));
verify(mockDataFileWriter2).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome2.getMessage())));
verify(mockDataFileWriter1).create(((IndexedRecord) ome.getMessage()).getSchema(), mockAzureBlobOutputStream1);
verify(mockDataFileWriter2).create(((IndexedRecord) ome2.getMessage()).getSchema(), mockAzureBlobOutputStream2);
}
use of com.azure.storage.blob.BlobContainerAsyncClient in project ambry by linkedin.
the class StorageClient method getAsyncContainerClient.
/**
* Get a reference to an Azure container async client, creating it if necessary.
* @param containerName the container name.
* @param autoCreate flag indicating whether to create the container if it does not exist.
* @return the created {@link BlobContainerClient}.
*/
private BlobContainerAsyncClient getAsyncContainerClient(String containerName, boolean autoCreate) throws ExecutionException, InterruptedException {
BlobContainerAsyncClient containerAsyncClient = storageAsyncClientRef.get().getBlobContainerAsyncClient(containerName);
if (autoCreate) {
if (!knownContainers.contains(containerName)) {
try {
if (!containerAsyncClient.exists().toFuture().get()) {
containerAsyncClient.create().block();
logger.info("Created container {}", containerName);
}
} catch (BlobStorageException ex) {
if (ex.getErrorCode() != BlobErrorCode.CONTAINER_ALREADY_EXISTS) {
logger.error("Failed to create container {}", containerName);
throw ex;
}
}
knownContainers.add(containerName);
}
}
return containerAsyncClient;
}
use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.
the class TestAzureBlobAvroWriter method setup.
@Before
public void setup() throws Exception {
threadPool = new ThreadPoolExecutor(1, 1, 60, TimeUnit.SECONDS, new LinkedBlockingDeque<>());
ome = createOME("Topic1");
encodedRecord = new byte[100];
BlobContainerAsyncClient mockContainerAsyncClient = PowerMockito.mock(BlobContainerAsyncClient.class);
mockDataFileWriter = mock(DataFileWriter.class);
mockAzureBlobOutputStream = mock(AzureBlobOutputStream.class);
mockBlockBlobAsyncClient = PowerMockito.mock(BlockBlobAsyncClient.class);
when(mockBlockBlobAsyncClient.getBlobUrl()).thenReturn("https://samza.blob.core.windows.net/fake-blob-url");
mockCompression = CompressionFactory.getInstance().getCompression(CompressionType.GZIP);
azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerAsyncClient, mock(AzureBlobWriterMetrics.class), threadPool, THRESHOLD, 60000, "test", mockDataFileWriter, mockAzureBlobOutputStream, mockBlockBlobAsyncClient, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, Long.MAX_VALUE, Long.MAX_VALUE, mockCompression, // keeping blob size and number of records unlimited
false));
doReturn(encodedRecord).when(azureBlobAvroWriter).encodeRecord((IndexedRecord) ome.getMessage());
}
use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.
the class TestAzureBlobAvroWriter method testMultipleBlobClose.
@Test
public void testMultipleBlobClose() throws Exception {
String blobUrlPrefix = "test";
long maxBlobSize = AzureBlobAvroWriter.DATAFILEWRITER_OVERHEAD + 1000;
long maxRecordsPerBlob = 10;
BlobContainerAsyncClient mockContainerClient = PowerMockito.mock(BlobContainerAsyncClient.class);
azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerClient, mock(AzureBlobWriterMetrics.class), threadPool, THRESHOLD, 60000, blobUrlPrefix, mockDataFileWriter, mockAzureBlobOutputStream, mockBlockBlobAsyncClient, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, maxBlobSize, maxRecordsPerBlob, mockCompression, false));
DataFileWriter<IndexedRecord> mockDataFileWriter2 = mock(DataFileWriter.class);
AzureBlobOutputStream mockAzureBlobOutputStream2 = mock(AzureBlobOutputStream.class);
when(mockAzureBlobOutputStream.getSize()).thenReturn((long) 1);
BlobAsyncClient mockBlobAsyncClient = mock(BlobAsyncClient.class);
doReturn(mockBlobAsyncClient).when(mockContainerClient).getBlobAsyncClient(anyString());
doReturn(mockBlockBlobAsyncClient).when(mockBlobAsyncClient).getBlockBlobAsyncClient();
PowerMockito.whenNew(AzureBlobOutputStream.class).withAnyArguments().thenReturn(mockAzureBlobOutputStream2);
PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter2);
for (int i = 0; i <= maxRecordsPerBlob; i++) {
azureBlobAvroWriter.write(ome);
}
// first OME creates the first blob and 11th OME creates the second blob.
azureBlobAvroWriter.close();
verify(mockDataFileWriter).close();
verify(mockDataFileWriter2).close();
}
Aggregations