def test_check_done_when_not_done(self):
tracker = OffsetRestrictionTracker(100, 200)
self.assertTrue(tracker.try_claim(150))
self.assertTrue(tracker.try_claim(175))
with self.assertRaises(ValueError):
tracker.check_done()
def test_try_claim(self): tracker = OffsetRestrictionTracker(100, 200) self.assertEqual((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_regular(self): tracker = OffsetRestrictionTracker(100, 200) self.assertTrue(tracker.try_claim(105)) self.assertTrue(tracker.try_claim(110)) checkpoint = tracker.checkpoint() self.assertEqual((100, 111), tracker.current_restriction()) self.assertEqual((111, 200), checkpoint)
def test_window_observing_checkpoint_on_first_window_after_prior_split(
self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(30)
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.0,
0,
2, # stop index < len(windows) representing a prior split had occurred
self.windowed_value,
self.restriction,
self.watermark_estimator_state,
self.restriction_provider,
restriction_tracker,
self.watermark_estimator)
expected_primary_split, expected_residual_split = (
self.create_split_in_window(31, (self.window1, )))
_, expected_residual_windows = (self.create_split_across_windows(
None, (self.window2, )))
hc.assert_that(primaries,
hc.contains_inanyorder(expected_primary_split))
hc.assert_that(
residuals,
hc.contains_inanyorder(
expected_residual_split,
expected_residual_windows,
))
self.assertEqual(stop_index, 1)
def test_window_observing_split_on_middle_window(self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(30)
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.2, 1, 3, self.windowed_value, self.restriction,
self.watermark_estimator_state, self.restriction_provider,
restriction_tracker, self.watermark_estimator)
# 20% of 1.7 windows = 20% of 170 offset left = 34 offset
# 30 + 34 = 64 split offset
expected_primary_split, expected_residual_split = (
self.create_split_in_window(64, (self.window2, )))
expected_primary_windows, expected_residual_windows = (
self.create_split_across_windows((self.window1, ),
(self.window3, )))
hc.assert_that(
primaries,
hc.contains_inanyorder(
expected_primary_split,
expected_primary_windows,
))
hc.assert_that(
residuals,
hc.contains_inanyorder(
expected_residual_split,
expected_residual_windows,
))
self.assertEqual(stop_index, 2)
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_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_non_window_observing_split_when_restriction_is_done(self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(100)
self.assertIsNone(
PerWindowInvoker._try_split(0.1, None, None, self.windowed_value,
self.restriction,
self.watermark_estimator_state,
self.restriction_provider,
restriction_tracker,
self.watermark_estimator))
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_non_window_observing_split(self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(30)
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.1, None, None, self.windowed_value, self.restriction,
self.watermark_estimator_state, self.restriction_provider,
restriction_tracker, self.watermark_estimator)
expected_primary_split, expected_residual_split = (
self.create_split_in_window(37, self.windowed_value.windows))
self.assertEqual([expected_primary_split], primaries)
self.assertEqual([expected_residual_split], residuals)
# We don't expect the stop index to be set for non window observing splits
self.assertIsNone(stop_index)
def test_self_checkpoint_immediately(self): restriction_tracker = OffsetRestrictionTracker(OffsetRange(0, 10)) threadsafe_tracker = ThreadsafeRestrictionTracker(restriction_tracker) threadsafe_tracker.defer_remainder() deferred_residual, deferred_time = threadsafe_tracker.deferred_status() expected_residual = OffsetRange(0, 10) self.assertEqual(deferred_residual, expected_residual) self.assertTrue(isinstance(deferred_time, timestamp.Duration)) self.assertEqual(deferred_time, 0)
def test_self_checkpoint_with_relative_time(self):
threadsafe_tracker = iobase.ThreadsafeRestrictionTracker(
OffsetRestrictionTracker(OffsetRange(0, 10)))
threadsafe_tracker.defer_remainder(timestamp.Duration(100))
time.sleep(2)
_, deferred_time = threadsafe_tracker.deferred_status()
self.assertTrue(isinstance(deferred_time, timestamp.Duration))
# The expectation = 100 - 2 - some_delta
self.assertTrue(deferred_time <= 98)
def test_window_observing_split_on_window_boundary_round_down(self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(30)
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.3, 0, 3, self.windowed_value, self.restriction,
self.watermark_estimator_state, self.restriction_provider,
restriction_tracker, self.watermark_estimator)
# 30% of 2.7 windows = 30% of 270 offset left = 81 offset
# 30 + 81 = 111 offset --> round to end of window 1
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, 1)
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_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_checkpoint_after_failed_claim(self):
tracker = OffsetRestrictionTracker(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))
checkpoint = tracker.checkpoint()
self.assertTrue((100, 161), tracker.current_restriction())
self.assertTrue((161, 200), checkpoint)
def test_window_observing_split_on_window_boundary_round_down_on_last_window(
self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(30)
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.9, 0, 3, self.windowed_value, self.restriction,
self.watermark_estimator_state, self.restriction_provider,
restriction_tracker, self.watermark_estimator)
# 90% of 2.7 windows = 90% of 270 offset left = 243 offset
# 30 + 243 = 273 offset --> prefer a split so round to end of window 2
# instead of no split
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_non_expose_apis(self):
threadsafe_tracker = iobase.ThreadsafeRestrictionTracker(
OffsetRestrictionTracker(OffsetRange(0, 10)))
tracker_view = iobase.RestrictionTrackerView(threadsafe_tracker)
with self.assertRaises(AttributeError):
tracker_view.check_done()
with self.assertRaises(AttributeError):
tracker_view.current_progress()
with self.assertRaises(AttributeError):
tracker_view.try_split()
with self.assertRaises(AttributeError):
tracker_view.deferred_status()
def test_api_expose(self):
threadsafe_tracker = iobase.ThreadsafeRestrictionTracker(
OffsetRestrictionTracker(OffsetRange(0, 10)))
tracker_view = iobase.RestrictionTrackerView(threadsafe_tracker)
current_restriction = tracker_view.current_restriction()
self.assertEqual(current_restriction, OffsetRange(0, 10))
self.assertTrue(tracker_view.try_claim(0))
tracker_view.defer_remainder()
deferred_remainder, deferred_watermark = (
threadsafe_tracker.deferred_status())
self.assertEqual(deferred_remainder, OffsetRange(1, 10))
self.assertEqual(deferred_watermark, timestamp.Duration())
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_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_self_checkpoint_with_absolute_time(self):
threadsafe_tracker = iobase.ThreadsafeRestrictionTracker(
OffsetRestrictionTracker(OffsetRange(0, 10)))
now = timestamp.Timestamp.now()
schedule_time = now + timestamp.Duration(100)
self.assertTrue(isinstance(schedule_time, timestamp.Timestamp))
threadsafe_tracker.defer_remainder(schedule_time)
time.sleep(2)
_, deferred_time = threadsafe_tracker.deferred_status()
self.assertTrue(isinstance(deferred_time, timestamp.Duration))
# The expectation =
# schedule_time - the time when deferred_status is called - some_delta
self.assertTrue(deferred_time <= 98)
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(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))
checkpoint = tracker.checkpoint()
self.assertTrue((100, 161), tracker.current_restriction())
self.assertTrue((161, 200), checkpoint)
def create_tracker(self, restriction):
return OffsetRestrictionTracker(restriction[0], restriction[1])
def test_claim_before_starting_range(self):
with self.assertRaises(ValueError):
tracker = OffsetRestrictionTracker(100, 200)
tracker.try_claim(90)
def create_tracker(self, restriction):
if self._disable_liquid_sharding:
return NonLiquidShardingOffsetRangeTracker(restriction)
else:
return OffsetRestrictionTracker(restriction)
def test_initialization(self):
with self.assertRaises(ValueError):
iobase.RestrictionTrackerView(
OffsetRestrictionTracker(OffsetRange(0, 10)))
def test_check_done_after_try_claim_past_end_of_range(self): tracker = OffsetRestrictionTracker(100, 200) self.assertTrue(tracker.try_claim(150)) self.assertTrue(tracker.try_claim(175)) self.assertFalse(tracker.try_claim(200)) tracker.check_done()
def test_checkpoint_unstarted(self): tracker = OffsetRestrictionTracker(100, 200) checkpoint = tracker.checkpoint() self.assertEqual((100, 100), tracker.current_restriction()) self.assertEqual((100, 200), checkpoint)
def test_check_done_after_try_claim_right_before_end_of_range(self): tracker = OffsetRestrictionTracker(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_checkpoint_just_started(self): tracker = OffsetRestrictionTracker(100, 200) self.assertTrue(tracker.try_claim(100)) checkpoint = tracker.checkpoint() self.assertEqual((100, 101), tracker.current_restriction()) self.assertEqual((101, 200), checkpoint)
def test_defer_remainder_with_wrong_time_type(self):
threadsafe_tracker = iobase.ThreadsafeRestrictionTracker(
OffsetRestrictionTracker(OffsetRange(0, 10)))
with self.assertRaises(ValueError):
threadsafe_tracker.defer_remainder(10)
def test_non_monotonic_claim(self):
with self.assertRaises(ValueError):
tracker = OffsetRestrictionTracker(100, 200)
self.assertTrue(tracker.try_claim(105))
self.assertTrue(tracker.try_claim(110))
self.assertTrue(tracker.try_claim(103))
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_claim_empty_range(self): tracker = OffsetRestrictionTracker(OffsetRange(0, 0)) self.assertFalse(tracker.try_claim(0))
def test_check_done_empty_range(self): tracker = OffsetRestrictionTracker(OffsetRange(0, 0)) tracker.check_done()
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_with_no_claims(self):
tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
with self.assertRaises(ValueError):
tracker.check_done()
