def test_window_observing_split_on_first_window(self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(30)
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.2, 0, 3, self.windowed_value, self.restriction,
self.watermark_estimator_state, self.restriction_provider,
restriction_tracker, self.watermark_estimator)
# 20% of 2.7 windows = 20% of 270 offset left = 54 offset
# 30 + 54 = 84 split offset
expected_primary_split, expected_residual_split = (
self.create_split_in_window(84, (self.window1, )))
_, expected_residual_windows = (self.create_split_across_windows(
None, (
self.window2,
self.window3,
)))
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_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_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_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_window_observing_split_on_last_window(self):
restriction_tracker = OffsetRestrictionTracker(self.restriction)
restriction_tracker.try_claim(30)
(primaries, residuals, stop_index) = PerWindowInvoker._try_split(
0.2, 2, 3, self.windowed_value, self.restriction,
self.watermark_estimator_state, self.restriction_provider,
restriction_tracker, self.watermark_estimator)
# 20% of 0.7 windows = 20% of 70 offset left = 14 offset
# 30 + 14 = 44 split offset
expected_primary_split, expected_residual_split = (
self.create_split_in_window(44, (self.window3, )))
expected_primary_windows, _ = (self.create_split_across_windows((
self.window1,
self.window2,
), None))
hc.assert_that(
primaries,
hc.contains_inanyorder(
expected_primary_split,
expected_primary_windows,
))
hc.assert_that(residuals,
hc.contains_inanyorder(expected_residual_split, ))
self.assertEqual(stop_index, 3)
