content: /*
   * 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.beam.sdk.io.range;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.rules.ExpectedException;
  import org.junit.runner.RunWith;
  import org.junit.runners.JUnit4;
  
  /** Tests for {@link OffsetRangeTracker}. */
  @RunWith(JUnit4.class)
  public class OffsetRangeTrackerTest {
    @Rule public final ExpectedException expected = ExpectedException.none();
  
    @Test
    public void testUpdateStartOffset() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      assertEquals(100, tracker.getStartPosition().longValue());
      // Update start offset to first record returned
      assertTrue(tracker.tryReturnRecordAt(true, 150));
      assertEquals(150, tracker.getStartPosition().longValue());
      assertTrue(tracker.tryReturnRecordAt(true, 180));
      assertEquals(150, tracker.getStartPosition().longValue());
    }
  
    @Test
    public void testTryReturnRecordSimpleSparse() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      assertTrue(tracker.tryReturnRecordAt(true, 110));
      assertTrue(tracker.tryReturnRecordAt(true, 140));
      assertTrue(tracker.tryReturnRecordAt(true, 183));
      assertFalse(tracker.tryReturnRecordAt(true, 210));
    }
  
    @Test
    public void testTryReturnRecordSimpleDense() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(3, 6);
      assertTrue(tracker.tryReturnRecordAt(true, 3));
      assertTrue(tracker.tryReturnRecordAt(true, 4));
      assertTrue(tracker.tryReturnRecordAt(true, 5));
      assertFalse(tracker.tryReturnRecordAt(true, 6));
    }
  
    @Test
    public void testTryReturnRecordContinuesUntilSplitPoint() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(9, 18);
      // Return records with gaps of 2; every 3rd record is a split point.
      assertTrue(tracker.tryReturnRecordAt(true, 10));
      assertTrue(tracker.tryReturnRecordAt(false, 12));
      assertTrue(tracker.tryReturnRecordAt(false, 14));
      assertTrue(tracker.tryReturnRecordAt(true, 16));
      // Out of range, but not a split point...
      assertTrue(tracker.tryReturnRecordAt(false, 18));
      assertTrue(tracker.tryReturnRecordAt(false, 20));
      // Out of range AND a split point.
      assertFalse(tracker.tryReturnRecordAt(true, 22));
    }
  
    @Test
    public void testSplitAtOffsetFailsIfUnstarted() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      assertFalse(tracker.trySplitAtPosition(150));
    }
  
    @Test
    public void testSplitAtOffset() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      assertTrue(tracker.tryReturnRecordAt(true, 110));
      // Example positions we shouldn't split at, when last record is [110, 130]:
      assertFalse(tracker.trySplitAtPosition(109));
      assertFalse(tracker.trySplitAtPosition(110));
      assertFalse(tracker.trySplitAtPosition(200));
      assertFalse(tracker.trySplitAtPosition(210));
      // Example positions we *should* split at:
      assertTrue(tracker.copy().trySplitAtPosition(111));
      assertTrue(tracker.copy().trySplitAtPosition(129));
      assertTrue(tracker.copy().trySplitAtPosition(130));
      assertTrue(tracker.copy().trySplitAtPosition(131));
      assertTrue(tracker.copy().trySplitAtPosition(150));
      assertTrue(tracker.copy().trySplitAtPosition(199));
  
      // If we split at 170 and then at 150:
      assertTrue(tracker.trySplitAtPosition(170));
      assertTrue(tracker.trySplitAtPosition(150));
      // Should be able to return a record starting before the new stop offset.
      // Returning records starting at the same offset is ok.
      assertTrue(tracker.copy().tryReturnRecordAt(true, 135));
      assertTrue(tracker.copy().tryReturnRecordAt(true, 135));
      // Should be able to return a record starting right before the new stop offset.
      assertTrue(tracker.copy().tryReturnRecordAt(true, 149));
      // Should not be able to return a record starting at or after the new stop offset
      assertFalse(tracker.tryReturnRecordAt(true, 150));
      assertFalse(tracker.tryReturnRecordAt(true, 151));
      // Should accept non-splitpoint records starting after stop offset.
      assertTrue(tracker.tryReturnRecordAt(false, 152));
      assertTrue(tracker.tryReturnRecordAt(false, 160));
      assertFalse(tracker.tryReturnRecordAt(true, 171));
    }
  
    @Test
    public void testGetPositionForFractionDense() throws Exception {
      // Represents positions 3, 4, 5.
      OffsetRangeTracker tracker = new OffsetRangeTracker(3, 6);
      // [3, 3) represents from [0, 1/3) fraction of [3, 6)
      assertEquals(3, tracker.getPositionForFractionConsumed(0.0));
      assertEquals(3, tracker.getPositionForFractionConsumed(1.0 / 6));
      assertEquals(3, tracker.getPositionForFractionConsumed(0.333));
      // [3, 4) represents from [0, 2/3) fraction of [3, 6)
      assertEquals(4, tracker.getPositionForFractionConsumed(0.334));
      assertEquals(4, tracker.getPositionForFractionConsumed(0.666));
      // [3, 5) represents from [0, 1) fraction of [3, 6)
      assertEquals(5, tracker.getPositionForFractionConsumed(0.667));
      assertEquals(5, tracker.getPositionForFractionConsumed(0.999));
      // The whole [3, 6) is consumed for fraction 1
      assertEquals(6, tracker.getPositionForFractionConsumed(1.0));
    }
  
    @Test
    public void testGetPositionForFractionDenseUpdateStartOffset() throws Exception {
      // Represents positions 3, 4, 5.
      OffsetRangeTracker tracker = new OffsetRangeTracker(3, 6);
      // [3, 3) represents from [0, 1/3) fraction of [3, 6)
      assertEquals(3, tracker.getPositionForFractionConsumed(0.333));
      // Update start offset to 4
      assertTrue(tracker.tryReturnRecordAt(true, 4));
      // [4, 4) represents from [0, 1/2) fraction of [4, 6)
      assertEquals(4, tracker.getPositionForFractionConsumed(0.0));
      assertEquals(4, tracker.getPositionForFractionConsumed(0.499));
      // [4, 5) represents from [0, 1) fraction of [4, 6)
      assertEquals(5, tracker.getPositionForFractionConsumed(0.5));
      assertEquals(5, tracker.getPositionForFractionConsumed(0.999));
      // The whole [4, 6) is consumed for fraction 1
      assertEquals(6, tracker.getPositionForFractionConsumed(1.0));
    }
  
    @Test
    public void testGetFractionConsumedDense() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(3, 6);
      assertEquals(0.0, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 3));
      assertEquals(0.0, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 4));
      assertEquals(1.0 / 3, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 5));
      assertEquals(2.0 / 3, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(false /* non-split-point */, 6));
      assertEquals(1.0, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(false /* non-split-point */, 7));
      assertEquals(1.0, tracker.getFractionConsumed(), 1e-6);
      assertFalse(tracker.tryReturnRecordAt(true, 7));
    }
  
    @Test
    public void testGetFractionConsumedSparse() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      assertEquals(0.0, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 100));
      assertEquals(0.0, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 110));
      // Consumed positions through 109 = total 10 positions of 100.
      assertEquals(0.1, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 150));
      assertEquals(0.5, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 195));
      assertEquals(0.95, tracker.getFractionConsumed(), 1e-6);
      assertFalse(tracker.tryReturnRecordAt(true, 200));
      assertEquals(1.0, tracker.getFractionConsumed(), 1e-6);
    }
  
    @Test
    public void testGetFractionConsumedUpdateStartOffset() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      assertTrue(tracker.tryReturnRecordAt(true, 150));
      assertEquals(0.0, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 160));
      assertEquals(0.2, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 180));
      assertEquals(0.6, tracker.getFractionConsumed(), 1e-6);
      assertTrue(tracker.tryReturnRecordAt(true, 195));
      assertEquals(0.9, tracker.getFractionConsumed(), 1e-6);
      assertFalse(tracker.tryReturnRecordAt(true, 200));
      assertEquals(1.0, tracker.getFractionConsumed(), 1e-6);
    }
  
    @Test
    public void testEverythingWithUnboundedRange() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, Long.MAX_VALUE);
      assertTrue(tracker.tryReturnRecordAt(true, 150));
      assertTrue(tracker.tryReturnRecordAt(true, 250));
      assertEquals(0.0, tracker.getFractionConsumed(), 1e-6);
      assertFalse(tracker.trySplitAtPosition(1000));
      try {
        tracker.getPositionForFractionConsumed(0.5);
        fail("getPositionForFractionConsumed should fail for an unbounded range");
      } catch (IllegalArgumentException e) {
        // Expected.
      }
    }
  
    @Test
    public void testTryReturnFirstRecordNotSplitPoint() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      expected.expect(IllegalStateException.class);
      tracker.tryReturnRecordAt(false, 120);
    }
  
    @Test
    public void testTryReturnRecordNonMonotonic() throws Exception {
      OffsetRangeTracker tracker = new OffsetRangeTracker(100, 200);
      tracker.tryReturnRecordAt(true, 120);
      expected.expect(IllegalStateException.class);
      tracker.tryReturnRecordAt(true, 110);
    }
  }
  