AzureBlobAvroWriter implements AzureBlobWriter { @Override public void write(OutgoingMessageEnvelope ome) throws IOException { Optional<IndexedRecord> optionalIndexedRecord; byte[] encodedRecord; if (ome.getMessage() instanceof IndexedRecord) { optionalIndexedRecord = Optional.of((IndexedRecord) ome.getMessage()); encodedRecord = encodeRecord((IndexedRecord) ome.getMessage()); } else if (ome.getMessage() instanceof byte[]) { optionalIndexedRecord = Optional.empty(); encodedRecord = (byte[]) ome.getMessage(); } else { throw new IllegalArgumentException("AzureBlobAvroWriter only supports IndexedRecord and byte[]."); } synchronized (currentDataFileWriterLock) { if (currentBlobWriterComponents == null || willCurrentBlobExceedSize(encodedRecord) || willCurrentBlobExceedRecordLimit()) { startNextBlob(optionalIndexedRecord); } currentBlobWriterComponents.dataFileWriter.appendEncoded(ByteBuffer.wrap(encodedRecord)); recordsInCurrentBlob++; currentBlobWriterComponents.azureBlobOutputStream.incrementNumberOfRecordsInBlob(); } } AzureBlobAvroWriter(BlobContainerAsyncClient containerAsyncClient, String blobURLPrefix,
      Executor blobThreadPool, AzureBlobWriterMetrics metrics,
      BlobMetadataGeneratorFactory blobMetadataGeneratorFactory, Config blobMetadataGeneratorConfig, String streamName,
      int maxBlockFlushThresholdSize, long flushTimeoutMs, Compression compression, boolean useRandomStringInBlobName,
      long maxBlobSize, long maxRecordsPerBlob); @VisibleForTesting  AzureBlobAvroWriter(BlobContainerAsyncClient containerAsyncClient, AzureBlobWriterMetrics metrics,
      Executor blobThreadPool, int maxBlockFlushThresholdSize, int flushTimeoutMs, String blobURLPrefix,
      DataFileWriter<IndexedRecord> dataFileWriter,
      AzureBlobOutputStream azureBlobOutputStream, BlockBlobAsyncClient blockBlobAsyncClient,
      BlobMetadataGeneratorFactory blobMetadataGeneratorFactory, Config blobMetadataGeneratorConfig, String streamName,
      long maxBlobSize, long maxRecordsPerBlob, Compression compression, boolean useRandomStringInBlobName); @Override void write(OutgoingMessageEnvelope ome); @Override void flush(); @Override void close();  }
@Test public void testWrite() throws Exception { int numberOfMessages = 10; for (int i = 0; i < numberOfMessages; ++i) { azureBlobAvroWriter.write(ome); } verify(mockDataFileWriter, times(numberOfMessages)).appendEncoded(ByteBuffer.wrap(encodedRecord)); verify(mockAzureBlobOutputStream, times(numberOfMessages)).incrementNumberOfRecordsInBlob(); }
@Test public void testWriteByteArray() throws Exception { OutgoingMessageEnvelope omeEncoded = new OutgoingMessageEnvelope(new SystemStream(SYSTEM_NAME, "Topic1"), "randomString".getBytes()); int numberOfMessages = 10; azureBlobAvroWriter.write(ome); for (int i = 0; i < numberOfMessages; ++i) { azureBlobAvroWriter.write(omeEncoded); } verify(mockDataFileWriter).appendEncoded(ByteBuffer.wrap(encodedRecord)); verify(mockDataFileWriter, times(numberOfMessages)).appendEncoded(ByteBuffer.wrap((byte[]) omeEncoded.getMessage())); verify(mockAzureBlobOutputStream, times(numberOfMessages + 1)).incrementNumberOfRecordsInBlob(); }
@Test(expected = IllegalStateException.class) public void testWriteByteArrayWithoutSchema() throws Exception { azureBlobAvroWriter = spy(new AzureBlobAvroWriter(PowerMockito.mock(BlobContainerAsyncClient.class), mock(AzureBlobWriterMetrics.class), threadPool, THRESHOLD, 60000, "test", null, null, null, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, 1000, 100, mockCompression, false)); OutgoingMessageEnvelope omeEncoded = new OutgoingMessageEnvelope(new SystemStream(SYSTEM_NAME, "Topic1"), new byte[100]); azureBlobAvroWriter.write(omeEncoded); }
@Test(expected = IOException.class) public void testWriteWhenDataFileWriterFails() throws Exception { doThrow(new IOException("Failed")).when(mockDataFileWriter).appendEncoded(ByteBuffer.wrap(encodedRecord)); azureBlobAvroWriter.write(ome); }