def test_server_watch(self):
with _test_server() as (server, port):
with grpc.insecure_channel('localhost:{}'.format(port)) as channel:
stub = key_value_pb2_grpc.KeyValueStoreStub(channel)
create_request = key_value_pb2.CreateRecordRequest(
record=key_value_pb2.Record(
name="golden-retriever",
value="pancakes",
))
stub.CreateRecord(create_request, wait_for_ready=True)
watch_request = key_value_pb2.WatchRecordRequest(
name="golden-retriever")
responses = queue.Queue()
def _gather_responses(stub, request):
response_iterator = stub.WatchRecord(request)
for response in response_iterator:
responses.put(response.value)
if response.value == _TEST_VALUES[-1]:
break
response_iterator.cancel()
thread_pool = concurrent.futures.ThreadPoolExecutor(
max_workers=1)
gather_future = thread_pool.submit(_gather_responses, stub,
watch_request)
_SENTINEL = "sentinel"
# Add meaningless values until the Watch connection has been confirmed
# to be established.
while responses.empty():
stub.UpdateRecord(
key_value_pb2.UpdateRecordRequest(
record=key_value_pb2.Record(
name="golden-retriever", value=_SENTINEL)))
# Append test values.
for value in _TEST_VALUES:
stub.UpdateRecord(
key_value_pb2.UpdateRecordRequest(
record=key_value_pb2.Record(
name="golden-retriever", value=value)))
# Ensure the queue has been filled with the collected values.
gather_future.result()
filtered_responses = [
response for response in responses.queue
if response != _SENTINEL
]
self.assertSequenceEqual(_TEST_VALUES, filtered_responses)
def _update(channel, key, value):
stub = key_value_pb2_grpc.KeyValueStoreStub(channel)
update_request = key_value_pb2.UpdateRecordRequest(
record=key_value_pb2.Record(
name=key,
value=value,
))
update_response = stub.UpdateRecord(update_request)
def _create(channel, key, value):
stub = key_value_pb2_grpc.KeyValueStoreStub(channel)
create_request = key_value_pb2.CreateRecordRequest(
record=key_value_pb2.Record(
name=key,
value=value,
))
record = stub.CreateRecord(create_request)
def GetRecord(self, request, context):
logging.info("Received Get request from {}".format(context.peer()))
if not self._kv_store.exists(request.name):
context.abort(
grpc.StatusCode.NOT_FOUND,
"Record at key '{}' does not exist.".format(request.name))
value = self._kv_store.get(request.name)
return key_value_pb2.Record(name=request.name, value=value)
def WatchRecord(self, request, context):
logging.info("Establishing Watch request with {}".format(
context.peer()))
if not self._kv_store.exists(request.name):
context.abort(
grpc.StatusCode.NOT_FOUND,
"Record at key '{}' does not exist.".format(request.name))
stop_event = threading.Event()
def on_rpc_done():
stop_event.set()
context.add_callback(on_rpc_done)
for value in self._kv_store.watch(request.name, stop_event):
yield key_value_pb2.Record(name=request.name, value=value)
logging.info("Terminated Watch request with {}".format(context.peer()))
def test_with_grpc_in_a_single_process(self):
with _test_server() as (server, port):
with grpc.insecure_channel('localhost:{}'.format(port)) as channel:
stub = key_value_pb2_grpc.KeyValueStoreStub(channel)
create_request = key_value_pb2.CreateRecordRequest(
record=key_value_pb2.Record(
name="golden-retriever",
value="pancakes",
))
create_response = stub.CreateRecord(create_request,
wait_for_ready=True)
self.assertEqual(create_request.record, create_response)
get_request = key_value_pb2.GetRecordRequest(
name="golden-retriever")
get_response = stub.GetRecord(get_request)
self.assertEqual(get_response, create_request.record)
def test_create_with_mock_server(self):
thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
fake_channel = grpc_testing.channel(
key_value_pb2.DESCRIPTOR.services_by_name.values(),
grpc_testing.strict_real_time())
# Offload the request to another thread so we can use this thread to
# fake the server-side results.
result_future = thread_pool.submit(
grpc_kv_client.create,
"localhost:1234",
"golden-retriever",
"pancakes",
channel=fake_channel)
# Fake the server-side results.
invocation_metadata, request, rpc = (fake_channel.take_unary_unary(
_get_method("CreateRecord")))
rpc.send_initial_metadata(())
rpc.terminate(
key_value_pb2.Record(name="golden-retriever", value="pancakes"),
(), grpc.StatusCode.OK, "")
# Ensure the client had the correct response.
result = result_future.result()
self.assertIsNone(result)