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)