AzureBlobOutputStream extends OutputStream { @Override public synchronized void close() { if (isClosed) { LOG.info("{}: already closed", blobAsyncClient.getBlobUrl().toString()); return; } LOG.info("{}: Close", blobAsyncClient.getBlobUrl().toString()); try { if (byteArrayOutputStream.isPresent()) { byteArrayOutputStream.get().close(); } if (blockList.size() == 0) { return; } CompletableFuture<Void> future = CompletableFuture.allOf(pendingUpload.toArray(new CompletableFuture[0])); LOG.info("Closing blob: {} PendingUpload:{} ", blobAsyncClient.getBlobUrl().toString(), pendingUpload.size()); future.get((long) flushTimeoutMs, TimeUnit.MILLISECONDS); LOG.info("For blob: {} committing blockList size:{}", blobAsyncClient.getBlobUrl().toString(), blockList.size()); metrics.updateAzureCommitMetrics(); BlobMetadataGenerator blobMetadataGenerator = getBlobMetadataGenerator(); commitBlob(blockList, blobMetadataGenerator.getBlobMetadata(new BlobMetadataContext(streamName, totalUploadedBlockSize, totalNumberOfRecordsInBlob))); } catch (Exception e) { String msg = String.format("Close blob %s failed with exception. Total pending sends %d", blobAsyncClient.getBlobUrl().toString(), pendingUpload.size()); throw new AzureException(msg, e); } finally { clearAndMarkClosed(); } } AzureBlobOutputStream(BlockBlobAsyncClient blobAsyncClient, Executor blobThreadPool, AzureBlobWriterMetrics metrics,
      BlobMetadataGeneratorFactory blobMetadataGeneratorFactory, Config blobMetadataGeneratorConfig, String streamName,
      long flushTimeoutMs, int maxBlockFlushThresholdSize, Compression compression); @VisibleForTesting  AzureBlobOutputStream(BlockBlobAsyncClient blobAsyncClient, Executor blobThreadPool, AzureBlobWriterMetrics metrics,
      BlobMetadataGeneratorFactory blobMetadataGeneratorFactory, Config blobMetadataGeneratorConfig, String streamName,
      long flushTimeoutMs, int maxBlockFlushThresholdSize,
      ByteArrayOutputStream byteArrayOutputStream, Compression compression); @Override synchronized void write(int b); @Override synchronized void write(byte[] b, int off, int len); @Override synchronized void flush(); @Override synchronized void close(); synchronized long getSize(); synchronized void releaseBuffer(); synchronized void incrementNumberOfRecordsInBlob();  }
@Test public void testClose() { azureBlobOutputStream.write(BYTES, 0, THRESHOLD); azureBlobOutputStream.incrementNumberOfRecordsInBlob(); int blockNum = 0; String blockId = String.format("%05d", blockNum); String blockIdEncoded = Base64.getEncoder().encodeToString(blockId.getBytes()); azureBlobOutputStream.close(); verify(mockMetrics).updateAzureCommitMetrics(); ArgumentCaptor<ArrayList> blockListArgument = ArgumentCaptor.forClass(ArrayList.class); ArgumentCaptor<Map> blobMetadataArg = ArgumentCaptor.forClass(Map.class); verify(azureBlobOutputStream).commitBlob(blockListArgument.capture(), blobMetadataArg.capture()); Assert.assertEquals(Arrays.asList(blockIdEncoded), blockListArgument.getAllValues().get(0)); Map<String, String> blobMetadata = (Map<String, String>) blobMetadataArg.getAllValues().get(0); Assert.assertEquals(blobMetadata.get(BLOB_RAW_SIZE_BYTES_METADATA), Long.toString(THRESHOLD)); Assert.assertEquals(blobMetadata.get(BLOB_STREAM_NAME_METADATA), FAKE_STREAM); Assert.assertEquals(blobMetadata.get(BLOB_RECORD_NUMBER_METADATA), Long.toString(1)); }