def test_examples_wordcount_streaming(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 FakeWriteStringsToPubSub(topic=None, values=None):
expected_topic = topic
def _inner(topic=None, subscription=None):
assert topic == expected_topic
return AssertTransform(equal_to(values))
return _inner
# Test basic execution.
input_topic = 'projects/fake-beam-test-project/topic/intopic'
input_values = [
TimestampedValue(b'a a b', 1),
TimestampedValue(u'🤷 ¯\\_(ツ)_/¯ b b '.encode('utf-8'), 12),
TimestampedValue(b'a b c c c', 20)
]
output_topic = 'projects/fake-beam-test-project/topic/outtopic'
output_values = ['a: 1', 'a: 2', 'b: 1', 'b: 3', 'c: 3']
beam.io.ReadFromPubSub = (
FakeReadFromPubSub(topic=input_topic, values=input_values))
beam.io.WriteStringsToPubSub = (
FakeWriteStringsToPubSub(topic=output_topic, values=output_values))
snippets.examples_wordcount_streaming([
'--input_topic',
'projects/fake-beam-test-project/topic/intopic',
'--output_topic',
'projects/fake-beam-test-project/topic/outtopic'
])
# Test with custom subscription.
input_sub = 'projects/fake-beam-test-project/subscriptions/insub'
beam.io.ReadFromPubSub = FakeReadFromPubSub(
subscription=input_sub, values=input_values)
snippets.examples_wordcount_streaming([
'--input_subscription',
'projects/fake-beam-test-project/subscriptions/insub',
'--output_topic',
'projects/fake-beam-test-project/topic/outtopic'
])
def test_examples_wordcount_streaming(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 FakeWriteStringsToPubSub(topic=None, values=None):
expected_topic = topic
def _inner(topic=None, subscription=None):
assert topic == expected_topic
return AssertTransform(equal_to(values))
return _inner
# Test basic execution.
input_topic = 'projects/fake-beam-test-project/topic/intopic'
input_values = [TimestampedValue('a a b', 1),
TimestampedValue(u'🤷 ¯\\_(ツ)_/¯ b b '.encode('utf-8'), 12),
TimestampedValue('a b c c c', 20)]
output_topic = 'projects/fake-beam-test-project/topic/outtopic'
output_values = ['a: 1', 'a: 2', 'b: 1', 'b: 3', 'c: 3']
beam.io.ReadFromPubSub = (
FakeReadFromPubSub(topic=input_topic, values=input_values))
beam.io.WriteStringsToPubSub = (
FakeWriteStringsToPubSub(topic=output_topic, values=output_values))
snippets.examples_wordcount_streaming([
'--input_topic', 'projects/fake-beam-test-project/topic/intopic',
'--output_topic', 'projects/fake-beam-test-project/topic/outtopic'])
# Test with custom subscription.
input_sub = 'projects/fake-beam-test-project/subscriptions/insub'
beam.io.ReadFromPubSub = FakeReadFromPubSub(subscription=input_sub,
values=input_values)
snippets.examples_wordcount_streaming([
'--input_subscription',
'projects/fake-beam-test-project/subscriptions/insub',
'--output_topic', 'projects/fake-beam-test-project/topic/outtopic'])