github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/restriction_trackers_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Unit tests for the range_trackers module."""
    19  
    20  # pytype: skip-file
    21  
    22  import logging
    23  import unittest
    24  
    25  from apache_beam.io.restriction_trackers import OffsetRange
    26  from apache_beam.io.restriction_trackers import OffsetRestrictionTracker
    27  
    28  
    29  class OffsetRangeTest(unittest.TestCase):
    30    def test_create(self):
    31      OffsetRange(0, 10)
    32      OffsetRange(10, 10)
    33      OffsetRange(10, 100)
    34  
    35      with self.assertRaises(ValueError):
    36        OffsetRange(10, 9)
    37  
    38    def test_split_respects_desired_num_splits(self):
    39      range = OffsetRange(10, 100)
    40      splits = list(range.split(desired_num_offsets_per_split=25))
    41      self.assertEqual(4, len(splits))
    42      self.assertIn(OffsetRange(10, 35), splits)
    43      self.assertIn(OffsetRange(35, 60), splits)
    44      self.assertIn(OffsetRange(60, 85), splits)
    45      self.assertIn(OffsetRange(85, 100), splits)
    46  
    47    def test_split_respects_min_num_splits(self):
    48      range = OffsetRange(10, 100)
    49      splits = list(
    50          range.split(
    51              desired_num_offsets_per_split=5, min_num_offsets_per_split=25))
    52      self.assertEqual(3, len(splits))
    53      self.assertIn(OffsetRange(10, 35), splits)
    54      self.assertIn(OffsetRange(35, 60), splits)
    55      self.assertIn(OffsetRange(60, 100), splits)
    56  
    57    def test_split_no_small_split_at_end(self):
    58      range = OffsetRange(10, 90)
    59      splits = list(range.split(desired_num_offsets_per_split=25))
    60      self.assertEqual(3, len(splits))
    61      self.assertIn(OffsetRange(10, 35), splits)
    62      self.assertIn(OffsetRange(35, 60), splits)
    63      self.assertIn(OffsetRange(60, 90), splits)
    64  
    65    def test_split_at(self):
    66      range = OffsetRange(0, 10)
    67      cur, residual = range.split_at(5)
    68      self.assertEqual(cur, OffsetRange(0, 5))
    69      self.assertEqual(residual, OffsetRange(5, 10))
    70  
    71  
    72  class OffsetRestrictionTrackerTest(unittest.TestCase):
    73    def test_try_claim(self):
    74      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
    75      self.assertEqual(OffsetRange(100, 200), tracker.current_restriction())
    76      self.assertTrue(tracker.try_claim(100))
    77      self.assertTrue(tracker.try_claim(150))
    78      self.assertTrue(tracker.try_claim(199))
    79      self.assertFalse(tracker.try_claim(200))
    80  
    81    def test_checkpoint_unstarted(self):
    82      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
    83      _, checkpoint = tracker.try_split(0)
    84      self.assertEqual(OffsetRange(100, 100), tracker.current_restriction())
    85      self.assertEqual(OffsetRange(100, 200), checkpoint)
    86  
    87    def test_checkpoint_just_started(self):
    88      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
    89      self.assertTrue(tracker.try_claim(100))
    90      _, checkpoint = tracker.try_split(0)
    91      self.assertEqual(OffsetRange(100, 101), tracker.current_restriction())
    92      self.assertEqual(OffsetRange(101, 200), checkpoint)
    93  
    94    def test_checkpoint_regular(self):
    95      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
    96      self.assertTrue(tracker.try_claim(105))
    97      self.assertTrue(tracker.try_claim(110))
    98      _, checkpoint = tracker.try_split(0)
    99      self.assertEqual(OffsetRange(100, 111), tracker.current_restriction())
   100      self.assertEqual(OffsetRange(111, 200), checkpoint)
   101  
   102    def test_checkpoint_claimed_last(self):
   103      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   104      self.assertTrue(tracker.try_claim(105))
   105      self.assertTrue(tracker.try_claim(110))
   106      self.assertTrue(tracker.try_claim(199))
   107      checkpoint = tracker.try_split(0)
   108      self.assertEqual(OffsetRange(100, 200), tracker.current_restriction())
   109      self.assertEqual(None, checkpoint)
   110  
   111    def test_checkpoint_after_failed_claim(self):
   112      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   113      self.assertTrue(tracker.try_claim(105))
   114      self.assertTrue(tracker.try_claim(110))
   115      self.assertTrue(tracker.try_claim(160))
   116      self.assertFalse(tracker.try_claim(240))
   117  
   118      self.assertIsNone(tracker.try_split(0))
   119      self.assertTrue(OffsetRange(100, 200), tracker.current_restriction())
   120  
   121    def test_non_monotonic_claim(self):
   122      with self.assertRaises(ValueError):
   123        tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   124        self.assertTrue(tracker.try_claim(105))
   125        self.assertTrue(tracker.try_claim(110))
   126        self.assertTrue(tracker.try_claim(103))
   127  
   128    def test_claim_before_starting_range(self):
   129      with self.assertRaises(ValueError):
   130        tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   131        tracker.try_claim(90)
   132  
   133    def test_check_done_after_try_claim_past_end_of_range(self):
   134      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   135      self.assertTrue(tracker.try_claim(150))
   136      self.assertTrue(tracker.try_claim(175))
   137      self.assertFalse(tracker.try_claim(220))
   138      tracker.check_done()
   139  
   140    def test_check_done_after_try_claim_right_before_end_of_range(self):
   141      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   142      self.assertTrue(tracker.try_claim(150))
   143      self.assertTrue(tracker.try_claim(175))
   144      self.assertTrue(tracker.try_claim(199))
   145      tracker.check_done()
   146  
   147    def test_check_done_when_not_done(self):
   148      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   149      self.assertTrue(tracker.try_claim(150))
   150      self.assertTrue(tracker.try_claim(175))
   151  
   152      with self.assertRaises(ValueError):
   153        tracker.check_done()
   154  
   155    def test_check_done_with_no_claims(self):
   156      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   157  
   158      with self.assertRaises(ValueError):
   159        tracker.check_done()
   160  
   161    def test_try_split(self):
   162      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   163      tracker.try_claim(100)
   164      cur, residual = tracker.try_split(0.5)
   165      self.assertEqual(OffsetRange(100, 150), cur)
   166      self.assertEqual(OffsetRange(150, 200), residual)
   167      self.assertEqual(cur, tracker.current_restriction())
   168  
   169    def test_try_split_when_restriction_is_done(self):
   170      tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
   171      tracker.try_claim(199)
   172      self.assertIsNone(tracker.try_split(0.5))
   173      tracker.try_claim(200)
   174      self.assertIsNone(tracker.try_split(0.5))
   175  
   176    def test_check_done_empty_range(self):
   177      tracker = OffsetRestrictionTracker(OffsetRange(0, 0))
   178      tracker.check_done()
   179  
   180    def test_try_claim_empty_range(self):
   181      tracker = OffsetRestrictionTracker(OffsetRange(0, 0))
   182      self.assertFalse(tracker.try_claim(0))
   183  
   184    def test_checkpoint_empty_range(self):
   185      tracker = OffsetRestrictionTracker(OffsetRange(0, 0))
   186      self.assertIsNone(tracker.try_split(0))
   187      self.assertFalse(tracker.try_claim(0))
   188      self.assertIsNone(tracker.try_split(0))
   189  
   190  
   191  if __name__ == '__main__':
   192    logging.getLogger().setLevel(logging.INFO)
   193    unittest.main()