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