def test_spec_construction(self):
BagStateSpec('statename', VarIntCoder())
with self.assertRaises(TypeError):
BagStateSpec(123, VarIntCoder())
CombiningValueStateSpec('statename', VarIntCoder(), TopCombineFn(10))
with self.assertRaises(TypeError):
CombiningValueStateSpec(123, VarIntCoder(), TopCombineFn(10))
with self.assertRaises(TypeError):
CombiningValueStateSpec('statename', VarIntCoder(), object())
SetStateSpec('setstatename', VarIntCoder())
with self.assertRaises(TypeError):
SetStateSpec(123, VarIntCoder())
with self.assertRaises(TypeError):
SetStateSpec('setstatename', object())
ReadModifyWriteStateSpec('valuestatename', VarIntCoder())
with self.assertRaises(TypeError):
ReadModifyWriteStateSpec(123, VarIntCoder())
with self.assertRaises(TypeError):
ReadModifyWriteStateSpec('valuestatename', object())
# TODO: add more spec tests
with self.assertRaises(ValueError):
DoFn.TimerParam(BagStateSpec('elements', BytesCoder()))
TimerSpec('timer', TimeDomain.WATERMARK)
TimerSpec('timer', TimeDomain.REAL_TIME)
with self.assertRaises(ValueError):
TimerSpec('timer', 'bogus_time_domain')
with self.assertRaises(ValueError):
DoFn.StateParam(TimerSpec('timer', TimeDomain.WATERMARK))
def test_state_spec_proto_conversion(self):
context = pipeline_context.PipelineContext()
state = BagStateSpec('statename', VarIntCoder())
state_proto = state.to_runner_api(context)
self.assertEqual(
beam_runner_api_pb2.FunctionSpec(
urn=common_urns.user_state.BAG.urn), state_proto.protocol)
context = pipeline_context.PipelineContext()
state = CombiningValueStateSpec('statename', VarIntCoder(),
TopCombineFn(10))
state_proto = state.to_runner_api(context)
self.assertEqual(
beam_runner_api_pb2.FunctionSpec(
urn=common_urns.user_state.BAG.urn), state_proto.protocol)
context = pipeline_context.PipelineContext()
state = SetStateSpec('setstatename', VarIntCoder())
state_proto = state.to_runner_api(context)
self.assertEqual(
beam_runner_api_pb2.FunctionSpec(
urn=common_urns.user_state.BAG.urn), state_proto.protocol)
context = pipeline_context.PipelineContext()
state = ReadModifyWriteStateSpec('valuestatename', VarIntCoder())
state_proto = state.to_runner_api(context)
self.assertEqual(
beam_runner_api_pb2.FunctionSpec(
urn=common_urns.user_state.BAG.urn), state_proto.protocol)