summaryrefslogtreecommitdiffstats
path: root/tests/mocks.py
blob: dc6e11720b91f09c93950bf68f1353511cc2e6a1 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from typing import List

from grpclib.client import Channel


class MockChannel(Channel):
    # noinspection PyMissingConstructor
    def __init__(self, responses=None) -> None:
        self.responses = responses or []
        self.requests = []
        self._loop = None

    def request(self, route, cardinality, request, response_type, **kwargs):
        self.requests.append(
            {
                "route": route,
                "cardinality": cardinality,
                "request": request,
                "response_type": response_type,
            }
        )
        return MockStream(self.responses)


class MockStream:
    def __init__(self, responses: List) -> None:
        super().__init__()
        self.responses = responses

    async def recv_message(self):
        return self.responses.pop(0)

    async def send_message(self, *args, **kwargs):
        pass

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return True

    async def __aenter__(self):
        return self