code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.buffered_stream_consumer;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  
  import com.fasterxml.jackson.databind.JsonNode;
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.protocol.models.AirbyteRecordMessage;
  import org.junit.jupiter.api.Test;
  
  class RecordSizeEstimatorTest {
  
    private static final JsonNode DATA_0 = Jsons.deserialize("{}");
    private static final JsonNode DATA_1 = Jsons.deserialize("{ \"field1\": true }");
    private static final JsonNode DATA_2 = Jsons.deserialize("{ \"field1\": 10000 }");
    private static final long DATA_0_SIZE = RecordSizeEstimator.getStringByteSize(DATA_0);
    private static final long DATA_1_SIZE = RecordSizeEstimator.getStringByteSize(DATA_1);
    private static final long DATA_2_SIZE = RecordSizeEstimator.getStringByteSize(DATA_2);
  
    @Test
    public void testPeriodicSampling() {
      // the estimate performs a size sampling every 3 records
      final RecordSizeEstimator sizeEstimator = new RecordSizeEstimator(3);
      final String stream = "stream";
      final AirbyteRecordMessage record0 = new AirbyteRecordMessage().withStream(stream).withData(DATA_0);
      final AirbyteRecordMessage record1 = new AirbyteRecordMessage().withStream(stream).withData(DATA_1);
      final AirbyteRecordMessage record2 = new AirbyteRecordMessage().withStream(stream).withData(DATA_2);
  
      // sample record message 1
      final long firstEstimation = DATA_1_SIZE;
      assertEquals(firstEstimation, sizeEstimator.getEstimatedByteSize(record1));
      // next two calls return the first sampling result
      assertEquals(firstEstimation, sizeEstimator.getEstimatedByteSize(record0));
      assertEquals(firstEstimation, sizeEstimator.getEstimatedByteSize(record0));
  
      // sample record message 2
      final long secondEstimation = firstEstimation / 2 + DATA_2_SIZE / 2;
      assertEquals(secondEstimation, sizeEstimator.getEstimatedByteSize(record2));
      // next two calls return the second sampling result
      assertEquals(secondEstimation, sizeEstimator.getEstimatedByteSize(record0));
      assertEquals(secondEstimation, sizeEstimator.getEstimatedByteSize(record0));
  
      // sample record message 1
      final long thirdEstimation = secondEstimation / 2 + DATA_1_SIZE / 2;
      assertEquals(thirdEstimation, sizeEstimator.getEstimatedByteSize(record1));
      // next two calls return the first sampling result
      assertEquals(thirdEstimation, sizeEstimator.getEstimatedByteSize(record0));
      assertEquals(thirdEstimation, sizeEstimator.getEstimatedByteSize(record0));
    }
  
    @Test
    public void testDifferentEstimationPerStream() {
      final RecordSizeEstimator sizeEstimator = new RecordSizeEstimator();
      final AirbyteRecordMessage record0 = new AirbyteRecordMessage().withStream("stream1").withData(DATA_0);
      final AirbyteRecordMessage record1 = new AirbyteRecordMessage().withStream("stream2").withData(DATA_1);
      final AirbyteRecordMessage record2 = new AirbyteRecordMessage().withStream("stream3").withData(DATA_2);
      assertEquals(DATA_0_SIZE, sizeEstimator.getEstimatedByteSize(record0));
      assertEquals(DATA_1_SIZE, sizeEstimator.getEstimatedByteSize(record1));
      assertEquals(DATA_2_SIZE, sizeEstimator.getEstimatedByteSize(record2));
    }
  
  }
  