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_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_window_observing_split_on_last_window_when_split_not_possible(
self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(100)
# We assume that we can't split this fully claimed restriction
self.assertIsNone(restriction_tracker.try_split(0))
self.assertIsNone(
PerWindowInvoker._try_split(0.0, 2, 3, self.windowed_value,
self.restriction,
self.watermark_estimator_state,
self.restriction_provider,
restriction_tracker,
self.watermark_estimator))
def test_window_observing_split_on_middle_window_fallback(self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(100)
# We assume that we can't split this fully claimed restriction
self.assertIsNone(restriction_tracker.try_split(0))
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.0, 1, 3, self.windowed_value, self.restriction,
self.watermark_estimator_state, self.restriction_provider,
restriction_tracker, self.watermark_estimator)
expected_primary_windows, expected_residual_windows = (
self.create_split_across_windows((
self.window1,
self.window2,
), (self.window3, )))
hc.assert_that(primaries,
hc.contains_inanyorder(expected_primary_windows, ))
hc.assert_that(residuals,
hc.contains_inanyorder(expected_residual_windows, ))
self.assertEqual(stop_index, 2)
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_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))
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))