def test_consume_end_of_stream(self):
stream = ReceiveStream(1)
stream.close()
message_type = Mock()
assert list(stream.consume(message_type)) == []
def test_consume_grpc_error(self):
stream = ReceiveStream(1)
error = GrpcError("boom", "details", "message")
stream.queue.put(error)
message_type = Mock()
with pytest.raises(GrpcError):
next(stream.consume(message_type))
def test_consume_uncompressed_message(self):
stream = ReceiveStream(1)
message_data = b"x"
message_type = Mock()
message = message_type()
stream.queue.put((False, message_data))
stream.close() # close stream so that consume exits
assert list(stream.consume(message_type)) == [message]
assert message.ParseFromString.call_args_list == [call(message_data)]