def test_multiple_sessions(self):
resp_a = self.stub.Events(beam_runner_api_pb2.EventsRequest())
resp_b = self.stub.Events(beam_runner_api_pb2.EventsRequest())
events_a = []
events_b = []
done = False
while not done:
a_is_done = False
b_is_done = False
try:
events_a.append(next(resp_a))
except StopIteration:
a_is_done = True
try:
events_b.append(next(resp_b))
except StopIteration:
b_is_done = True
done = a_is_done and b_is_done
expected_events = [e for e in self.events()]
self.assertEqual(events_a, expected_events)
self.assertEqual(events_b, expected_events)
def test_multiple_sessions(self):
resp_a = self.stub.Events(
beam_runner_api_pb2.EventsRequest(output_ids=EXPECTED_KEYS))
resp_b = self.stub.Events(
beam_runner_api_pb2.EventsRequest(output_ids=EXPECTED_KEYS))
events_a = []
events_b = []
done = False
while not done:
a_is_done = False
b_is_done = False
try:
events_a.append(next(resp_a))
except StopIteration:
a_is_done = True
try:
events_b.append(next(resp_b))
except StopIteration:
b_is_done = True
done = a_is_done and b_is_done
expected_events = [
e for e in EventsReader(
expected_key=[EXPECTED_KEYS]).read_multiple([EXPECTED_KEYS])
]
self.assertEqual(events_a, expected_events)
self.assertEqual(events_b, expected_events)
def _stream_events_from_rpc(endpoint, output_tags, coder, channel,
is_alive):
"""Yields the events received from the given endpoint.
This is the producer thread that reads events from the TestStreamService and
puts them onto the shared queue. At the end of the stream, an _EndOfStream
is placed on the channel to signify a successful end.
"""
stub_channel = grpc.insecure_channel(endpoint)
stub = beam_runner_api_pb2_grpc.TestStreamServiceStub(stub_channel)
# Request the PCollections that we are looking for from the service.
event_request = beam_runner_api_pb2.EventsRequest(
output_ids=[str(tag) for tag in output_tags])
event_stream = stub.Events(event_request)
try:
for e in event_stream:
channel.put(_TestStream.test_stream_payload_to_events(
e, coder))
if not is_alive():
return
except grpc.RpcError as e:
# Do not raise an exception in the non-error status codes. These can occur
# when the Python interpreter shuts down or when in a notebook environment
# when the kernel is interrupted.
if e.code() in (grpc.StatusCode.CANCELLED,
grpc.StatusCode.UNAVAILABLE):
return
raise e
finally:
# Gracefully stop the job if there is an exception.
channel.put(_EndOfStream())
def test_normal_run(self):
r = self.stub.Events(
beam_runner_api_pb2.EventsRequest(output_ids=EXPECTED_KEYS))
events = [e for e in r]
expected_events = [
e for e in EventsReader(
expected_key=[EXPECTED_KEYS]).read_multiple([EXPECTED_KEYS])
]
self.assertEqual(events, expected_events)
def events_from_rpc(endpoint, output_tags, coder):
"""Yields the events received from the given endpoint.
"""
stub_channel = grpc.insecure_channel(endpoint)
stub = beam_runner_api_pb2_grpc.TestStreamServiceStub(stub_channel)
# Request the PCollections that we are looking for from the service.
event_request = beam_runner_api_pb2.EventsRequest(
output_ids=[str(tag) for tag in output_tags])
event_stream = stub.Events(event_request)
for e in event_stream:
yield _TestStream.test_stream_payload_to_events(e, coder)
def _stream_events_from_rpc(endpoint, output_tags, coder, channel,
is_alive):
"""Yields the events received from the given endpoint.
This is the producer thread that reads events from the TestStreamService and
puts them onto the shared queue. At the end of the stream, an _EndOfStream
is placed on the channel to signify a successful end.
"""
stub_channel = grpc.insecure_channel(endpoint)
stub = beam_runner_api_pb2_grpc.TestStreamServiceStub(stub_channel)
# Request the PCollections that we are looking for from the service.
event_request = beam_runner_api_pb2.EventsRequest(
output_ids=[str(tag) for tag in output_tags])
event_stream = stub.Events(event_request)
for e in event_stream:
channel.put(_TestStream.test_stream_payload_to_events(e, coder))
if not is_alive():
return
channel.put(_EndOfStream())
def events_from_rpc(endpoint, output_tags, coder):
"""Yields the events received from the given endpoint.
"""
stub_channel = grpc.insecure_channel(endpoint)
stub = beam_runner_api_pb2_grpc.TestStreamServiceStub(stub_channel)
# Request the PCollections that we are looking for from the service.
event_request = beam_runner_api_pb2.EventsRequest(
output_ids=[str(tag) for tag in output_tags])
event_stream = stub.Events(event_request, timeout=30)
try:
while True:
yield _TestStream.test_stream_payload_to_events(
next(event_stream), coder)
except StopIteration:
return
except grpc.RpcError as e:
if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
_LOGGER.warning(
'TestStream timed out waiting for new events from service.'
' Stopping pipeline.')
return
raise e
def test_normal_run(self):
r = self.stub.Events(beam_runner_api_pb2.EventsRequest())
events = [e for e in r]
expected_events = [e for e in self.events()]
self.assertEqual(events, expected_events)
