tests/tests_aio/unit/_common.py

Go to the documentation of this file.

# Copyright 2020 The gRPC Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
 
import asyncio
from typing import AsyncIterable
 
import grpc
from grpc.aio._metadata import Metadata
from grpc.aio._typing import MetadataKey
from grpc.aio._typing import MetadataValue
from grpc.aio._typing import MetadatumType
from grpc.experimental import aio
 
from tests.unit.framework.common import test_constants
 
ADHOC_METHOD = '/test/AdHoc'
 
 
def seen_metadata(expected: Metadata, actual: Metadata):
    return not bool(set(tuple(expected)) - set(tuple(actual)))
 
 
def seen_metadatum(expected_key: MetadataKey, expected_value: MetadataValue,
                   actual: Metadata) -> bool:
    obtained = actual[expected_key]
    return obtained == expected_value
 
 
async def block_until_certain_state(channel: aio.Channel,
                                    expected_state: grpc.ChannelConnectivity):
    state = channel.get_state()
    while state != expected_state:
        await channel.wait_for_state_change(state)
        state = channel.get_state()
 
 
def inject_callbacks(call: aio.Call):
    first_callback_ran = asyncio.Event()
 
    def first_callback(call):
        # Validate that all resopnses have been received
        # and the call is an end state.
        assert call.done()
        first_callback_ran.set()
 
    second_callback_ran = asyncio.Event()
 
    def second_callback(call):
        # Validate that all responses have been received
        # and the call is an end state.
        assert call.done()
        second_callback_ran.set()
 
    call.add_done_callback(first_callback)
    call.add_done_callback(second_callback)
 
    async def validation():
        await asyncio.wait_for(
            asyncio.gather(first_callback_ran.wait(),
                           second_callback_ran.wait()),
            test_constants.SHORT_TIMEOUT)
 
    return validation()
 
 
class CountingRequestIterator:
 
    def __init__(self, request_iterator):
        self.request_cnt = 0
        self._request_iterator = request_iterator
 
    async def _forward_requests(self):
        async for request in self._request_iterator:
            self.request_cnt += 1
            yield request
 
    def __aiter__(self):
        return self._forward_requests()
 
 
class CountingResponseIterator:
 
    def __init__(self, response_iterator):
        self.response_cnt = 0
        self._response_iterator = response_iterator
 
    async def _forward_responses(self):
        async for response in self._response_iterator:
            self.response_cnt += 1
            yield response
 
    def __aiter__(self):
        return self._forward_responses()
 
 
class AdhocGenericHandler(grpc.GenericRpcHandler):
    """A generic handler to plugin testing server methods on the fly."""
    _handler: grpc.RpcMethodHandler
 
    def __init__(self):
        self._handler = None
 
    def set_adhoc_handler(self, handler: grpc.RpcMethodHandler):
        self._handler = handler
 
    def service(self, handler_call_details):
        if handler_call_details.method == ADHOC_METHOD:
            return self._handler
        else:
            return None