| # | |
| # 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() |