def test_streaming_wordcount_it(self):
# Build expected dataset.
expected_msg = [('%d: 1' % num).encode('utf-8')
for num in range(DEFAULT_INPUT_NUMBERS)]
# Set extra options to the pipeline for test purpose
state_verifier = PipelineStateMatcher(PipelineState.RUNNING)
pubsub_msg_verifier = PubSubMessageMatcher(self.project,
self.output_sub.name,
expected_msg,
timeout=400)
extra_opts = {
'input_subscription': self.input_sub.name,
'output_topic': self.output_topic.name,
'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION,
'on_success_matcher': all_of(state_verifier, pubsub_msg_verifier)
}
# Generate input data and inject to PubSub.
self._inject_numbers(self.input_topic, DEFAULT_INPUT_NUMBERS)
# Get pipeline options from command argument: --test-pipeline-options,
# and start pipeline job by calling pipeline main function.
streaming_wordcount_debugging.run(
self.test_pipeline.get_full_options_as_args(**extra_opts),
save_main_session=False)
def test_streaming_wordcount_debugging(self, *unused_mocks):
def FakeReadFromPubSub(topic=None, subscription=None, values=None):
expected_topic = topic
expected_subscription = subscription
def _inner(topic=None, subscription=None):
assert topic == expected_topic
assert subscription == expected_subscription
return TestStream().add_elements(values)
return _inner
class AssertTransform(beam.PTransform):
def __init__(self, matcher):
self.matcher = matcher
def expand(self, pcoll):
assert_that(pcoll, self.matcher)
def FakeWriteToPubSub(topic=None, values=None):
expected_topic = topic
def _inner(topic=None, subscription=None):
assert topic == expected_topic
return AssertTransform(equal_to(values))
return _inner
input_topic = 'projects/fake-beam-test-project/topic/intopic'
input_values = [
'150', '151', '152', '153', '154', '210', '211', '212', '213',
'214'
]
output_topic = 'projects/fake-beam-test-project/topic/outtopic'
output_values = [
'150: 1', '151: 1', '152: 1', '153: 1', '154: 1', '210: 1',
'211: 1', '212: 1', '213: 1', '214: 1'
]
beam.io.ReadFromPubSub = (FakeReadFromPubSub(
topic=input_topic,
values=list(x.encode('utf-8') for x in input_values)))
beam.io.WriteToPubSub = (FakeWriteToPubSub(
topic=output_topic,
values=list(x.encode('utf-8') for x in output_values)))
streaming_wordcount_debugging.run([
'--input_topic', 'projects/fake-beam-test-project/topic/intopic',
'--output_topic', 'projects/fake-beam-test-project/topic/outtopic'
],
save_main_session=False)
def test_streaming_wordcount_debugging_it(self):
# Set extra options to the pipeline for test purpose
state_verifier = PipelineStateMatcher(PipelineState.RUNNING)
pubsub_msg_verifier = PubSubMessageMatcher(
self.project, self.output_sub.name, EXPECTED_MESSAGE, timeout=400)
extra_opts = {
'input_subscription': self.input_sub.name,
'output_topic': self.output_topic.name,
'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION,
'on_success_matcher': all_of(state_verifier, pubsub_msg_verifier)
}
# Generate input data and inject to PubSub.
self._inject_data(self.input_topic, SAMPLE_MESSAGES)
# Get pipeline options from command argument: --test-pipeline-options,
# and start pipeline job by calling pipeline main function.
streaming_wordcount_debugging.run(
self.test_pipeline.get_full_options_as_args(**extra_opts),
save_main_session=False)