#

# 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.

#

"""Unit tests for the range_trackers module."""

# pytype: skip-file

import logging

import unittest

from apache_beam.io.restriction_trackers import OffsetRange

from apache_beam.io.restriction_trackers import OffsetRestrictionTracker

class OffsetRangeTest(unittest.TestCase):

def test_create(self):

OffsetRange(0, 10)

OffsetRange(10, 10)

OffsetRange(10, 100)

with self.assertRaises(ValueError):

OffsetRange(10, 9)

def test_split_respects_desired_num_splits(self):

range = OffsetRange(10, 100)

splits = list(range.split(desired_num_offsets_per_split=25))

self.assertEqual(4, len(splits))

self.assertIn(OffsetRange(10, 35), splits)

self.assertIn(OffsetRange(35, 60), splits)

self.assertIn(OffsetRange(60, 85), splits)

self.assertIn(OffsetRange(85, 100), splits)

def test_split_respects_min_num_splits(self):

range = OffsetRange(10, 100)

splits = list(

range.split(

desired_num_offsets_per_split=5, min_num_offsets_per_split=25))

self.assertEqual(3, len(splits))

self.assertIn(OffsetRange(10, 35), splits)

self.assertIn(OffsetRange(35, 60), splits)

self.assertIn(OffsetRange(60, 100), splits)

def test_split_no_small_split_at_end(self):

range = OffsetRange(10, 90)

splits = list(range.split(desired_num_offsets_per_split=25))

self.assertEqual(3, len(splits))

self.assertIn(OffsetRange(10, 35), splits)

self.assertIn(OffsetRange(35, 60), splits)

self.assertIn(OffsetRange(60, 90), splits)

def test_split_at(self):

range = OffsetRange(0, 10)

cur, residual = range.split_at(5)

self.assertEqual(cur, OffsetRange(0, 5))

self.assertEqual(residual, OffsetRange(5, 10))

class OffsetRestrictionTrackerTest(unittest.TestCase):

def test_try_claim(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertEqual(OffsetRange(100, 200), tracker.current_restriction())

self.assertTrue(tracker.try_claim(100))

self.assertTrue(tracker.try_claim(150))

self.assertTrue(tracker.try_claim(199))

self.assertFalse(tracker.try_claim(200))

def test_checkpoint_unstarted(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

_, checkpoint = tracker.try_split(0)

self.assertEqual(OffsetRange(100, 100), tracker.current_restriction())

self.assertEqual(OffsetRange(100, 200), checkpoint)

def test_checkpoint_just_started(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(100))

_, checkpoint = tracker.try_split(0)

self.assertEqual(OffsetRange(100, 101), tracker.current_restriction())

self.assertEqual(OffsetRange(101, 200), checkpoint)

def test_checkpoint_regular(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(105))

self.assertTrue(tracker.try_claim(110))

_, checkpoint = tracker.try_split(0)

self.assertEqual(OffsetRange(100, 111), tracker.current_restriction())

self.assertEqual(OffsetRange(111, 200), checkpoint)

def test_checkpoint_claimed_last(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(105))

self.assertTrue(tracker.try_claim(110))

self.assertTrue(tracker.try_claim(199))

checkpoint = tracker.try_split(0)

self.assertEqual(OffsetRange(100, 200), tracker.current_restriction())

self.assertEqual(None, checkpoint)

def test_checkpoint_after_failed_claim(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(105))

self.assertTrue(tracker.try_claim(110))

self.assertTrue(tracker.try_claim(160))

self.assertFalse(tracker.try_claim(240))

self.assertIsNone(tracker.try_split(0))

self.assertTrue(OffsetRange(100, 200), tracker.current_restriction())

def test_non_monotonic_claim(self):

with self.assertRaises(ValueError):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(105))

self.assertTrue(tracker.try_claim(110))

self.assertTrue(tracker.try_claim(103))

def test_claim_before_starting_range(self):

with self.assertRaises(ValueError):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

tracker.try_claim(90)

def test_check_done_after_try_claim_past_end_of_range(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(150))

self.assertTrue(tracker.try_claim(175))

self.assertFalse(tracker.try_claim(220))

tracker.check_done()

def test_check_done_after_try_claim_right_before_end_of_range(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(150))

self.assertTrue(tracker.try_claim(175))

self.assertTrue(tracker.try_claim(199))

tracker.check_done()

def test_check_done_when_not_done(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

self.assertTrue(tracker.try_claim(150))

self.assertTrue(tracker.try_claim(175))

with self.assertRaises(ValueError):

tracker.check_done()

def test_check_done_with_no_claims(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

with self.assertRaises(ValueError):

tracker.check_done()

def test_try_split(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

tracker.try_claim(100)

cur, residual = tracker.try_split(0.5)

self.assertEqual(OffsetRange(100, 150), cur)

self.assertEqual(OffsetRange(150, 200), residual)

self.assertEqual(cur, tracker.current_restriction())

def test_try_split_when_restriction_is_done(self):

tracker = OffsetRestrictionTracker(OffsetRange(100, 200))

tracker.try_claim(199)

self.assertIsNone(tracker.try_split(0.5))

tracker.try_claim(200)

self.assertIsNone(tracker.try_split(0.5))

def test_check_done_empty_range(self):

tracker = OffsetRestrictionTracker(OffsetRange(0, 0))

tracker.check_done()

def test_try_claim_empty_range(self):

tracker = OffsetRestrictionTracker(OffsetRange(0, 0))

self.assertFalse(tracker.try_claim(0))

def test_checkpoint_empty_range(self):

tracker = OffsetRestrictionTracker(OffsetRange(0, 0))

self.assertIsNone(tracker.try_split(0))

self.assertFalse(tracker.try_claim(0))

self.assertIsNone(tracker.try_split(0))

if __name__ == '__main__':

logging.getLogger().setLevel(logging.INFO)

unittest.main()