def _validate_stateful_dofn(self):
userstate.validate_stateful_dofn(self.do_fn)
def _validate_stateful_dofn(self): userstate.validate_stateful_dofn(self.do_fn)
def test_validation_typos(self):
# (1) Here, the user mistakenly used the same timer spec twice for two
# different timer callbacks.
with self.assertRaisesRegex(
ValueError,
r'Multiple on_timer callbacks registered for TimerSpec\(.*expiry1\).'
):
class StatefulDoFnWithTimerWithTypo1(DoFn): # pylint: disable=unused-variable
BUFFER_STATE = BagStateSpec('buffer', BytesCoder())
EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK)
EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK)
def process(self, element):
pass
@on_timer(EXPIRY_TIMER_1)
def on_expiry_1(self,
buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired1'
# Note that we mistakenly associate this with the first timer.
@on_timer(EXPIRY_TIMER_1)
def on_expiry_2(self,
buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired2'
# (2) Here, the user mistakenly used the same callback name and overwrote
# the first on_expiry_1 callback.
class StatefulDoFnWithTimerWithTypo2(DoFn):
BUFFER_STATE = BagStateSpec('buffer', BytesCoder())
EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK)
EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK)
def process(self,
element,
timer1=DoFn.TimerParam(EXPIRY_TIMER_1),
timer2=DoFn.TimerParam(EXPIRY_TIMER_2)):
pass
@on_timer(EXPIRY_TIMER_1)
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired1'
# Note that we mistakenly reuse the "on_expiry_1" name; this is valid
# syntactically in Python.
@on_timer(EXPIRY_TIMER_2)
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired2'
# Use a stable string value for matching.
def __repr__(self):
return 'StatefulDoFnWithTimerWithTypo2'
dofn = StatefulDoFnWithTimerWithTypo2()
with self.assertRaisesRegex(
ValueError,
(r'The on_timer callback for TimerSpec\(.*expiry1\) is not the '
r'specified .on_expiry_1 method for DoFn '
r'StatefulDoFnWithTimerWithTypo2 \(perhaps it was overwritten\?\).'
)):
validate_stateful_dofn(dofn)
# (2) Here, the user forgot to add an on_timer decorator for 'expiry2'
class StatefulDoFnWithTimerWithTypo3(DoFn):
BUFFER_STATE = BagStateSpec('buffer', BytesCoder())
EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK)
EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK)
def process(self,
element,
timer1=DoFn.TimerParam(EXPIRY_TIMER_1),
timer2=DoFn.TimerParam(EXPIRY_TIMER_2)):
pass
@on_timer(EXPIRY_TIMER_1)
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired1'
def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired2'
# Use a stable string value for matching.
def __repr__(self):
return 'StatefulDoFnWithTimerWithTypo3'
dofn = StatefulDoFnWithTimerWithTypo3()
with self.assertRaisesRegex(
ValueError,
(r'DoFn StatefulDoFnWithTimerWithTypo3 has a TimerSpec without an '
r'associated on_timer callback: TimerSpec\(.*expiry2\).')):
validate_stateful_dofn(dofn)
def test_validation_typos(self):
# (1) Here, the user mistakenly used the same timer spec twice for two
# different timer callbacks.
with self.assertRaisesRegexp(
ValueError,
r'Multiple on_timer callbacks registered for TimerSpec\(expiry1\).'):
class StatefulDoFnWithTimerWithTypo1(DoFn): # pylint: disable=unused-variable
BUFFER_STATE = BagStateSpec('buffer', BytesCoder())
EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK)
EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK)
def process(self, element):
pass
@on_timer(EXPIRY_TIMER_1)
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired1'
# Note that we mistakenly associate this with the first timer.
@on_timer(EXPIRY_TIMER_1)
def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired2'
# (2) Here, the user mistakenly used the same callback name and overwrote
# the first on_expiry_1 callback.
class StatefulDoFnWithTimerWithTypo2(DoFn):
BUFFER_STATE = BagStateSpec('buffer', BytesCoder())
EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK)
EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK)
def process(self, element,
timer1=DoFn.TimerParam(EXPIRY_TIMER_1),
timer2=DoFn.TimerParam(EXPIRY_TIMER_2)):
pass
@on_timer(EXPIRY_TIMER_1)
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired1'
# Note that we mistakenly reuse the "on_expiry_1" name; this is valid
# syntactically in Python.
@on_timer(EXPIRY_TIMER_2)
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired2'
# Use a stable string value for matching.
def __repr__(self):
return 'StatefulDoFnWithTimerWithTypo2'
dofn = StatefulDoFnWithTimerWithTypo2()
with self.assertRaisesRegexp(
ValueError,
(r'The on_timer callback for TimerSpec\(expiry1\) is not the '
r'specified .on_expiry_1 method for DoFn '
r'StatefulDoFnWithTimerWithTypo2 \(perhaps it was overwritten\?\).')):
validate_stateful_dofn(dofn)
# (2) Here, the user forgot to add an on_timer decorator for 'expiry2'
class StatefulDoFnWithTimerWithTypo3(DoFn):
BUFFER_STATE = BagStateSpec('buffer', BytesCoder())
EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK)
EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK)
def process(self, element,
timer1=DoFn.TimerParam(EXPIRY_TIMER_1),
timer2=DoFn.TimerParam(EXPIRY_TIMER_2)):
pass
@on_timer(EXPIRY_TIMER_1)
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired1'
def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired2'
# Use a stable string value for matching.
def __repr__(self):
return 'StatefulDoFnWithTimerWithTypo3'
dofn = StatefulDoFnWithTimerWithTypo3()
with self.assertRaisesRegexp(
ValueError,
(r'DoFn StatefulDoFnWithTimerWithTypo3 has a TimerSpec without an '
r'associated on_timer callback: TimerSpec\(expiry2\).')):
validate_stateful_dofn(dofn)
