Merge pull request #32 from nat-n/improve_stub

Add ability to provide metadata, timeout & deadline args to requests

betterproto/tests/service.proto

@@ -0,0 +1,15 @@
+syntax = "proto3";
+
+package service;
+
+message DoThingRequest {
+  int32 iterations = 1;
+}
+
+message DoThingResponse {
+  int32 successfulIterations = 1;
+}
+
+service ExampleService {
+  rpc DoThing (DoThingRequest) returns (DoThingResponse);
+}

betterproto/tests/test_service_stub.py

@@ -0,0 +1,123 @@
+import betterproto
+import grpclib
+from grpclib.testing import ChannelFor
+import pytest
+from typing import Dict
+from .service import DoThingResponse, DoThingRequest, ExampleServiceStub
+
+
+class ExampleService:
+    def __init__(self, test_hook=None):
+        # This lets us pass assertions to the servicer ;)
+        self.test_hook = test_hook
+
+    async def DoThing(
+        self, stream: "grpclib.server.Stream[DoThingRequest, DoThingResponse]"
+    ):
+        request = await stream.recv_message()
+        print("self.test_hook", self.test_hook)
+        if self.test_hook is not None:
+            self.test_hook(stream)
+        for iteration in range(request.iterations):
+            pass
+        await stream.send_message(DoThingResponse(request.iterations))
+
+    def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
+        return {
+            "/service.ExampleService/DoThing": grpclib.const.Handler(
+                self.DoThing,
+                grpclib.const.Cardinality.UNARY_UNARY,
+                DoThingRequest,
+                DoThingResponse,
+            ),
+        }
+
+
+async def _test_stub(stub, iterations=42, **kwargs):
+    response = await stub.do_thing(iterations=iterations)
+    assert response.successful_iterations == iterations
+
+
+def _get_server_side_test(deadline, metadata):
+    def server_side_test(stream):
+        assert stream.deadline._timestamp == pytest.approx(
+            deadline._timestamp, 1
+        ), "The provided deadline should be recieved serverside"
+        assert (
+            stream.metadata["authorization"] == metadata["authorization"]
+        ), "The provided authorization metadata should be recieved serverside"
+
+    return server_side_test
+
+
+@pytest.mark.asyncio
+async def test_simple_service_call():
+    async with ChannelFor([ExampleService()]) as channel:
+        await _test_stub(ExampleServiceStub(channel))
+
+
+@pytest.mark.asyncio
+async def test_service_call_with_upfront_request_params():
+    # Setting deadline
+    deadline = grpclib.metadata.Deadline.from_timeout(22)
+    metadata = {"authorization": "12345"}
+    async with ChannelFor(
+        [ExampleService(test_hook=_get_server_side_test(deadline, metadata))]
+    ) as channel:
+        await _test_stub(
+            ExampleServiceStub(channel, deadline=deadline, metadata=metadata)
+        )
+
+    # Setting timeout
+    timeout = 99
+    deadline = grpclib.metadata.Deadline.from_timeout(timeout)
+    metadata = {"authorization": "12345"}
+    async with ChannelFor(
+        [ExampleService(test_hook=_get_server_side_test(deadline, metadata))]
+    ) as channel:
+        await _test_stub(
+            ExampleServiceStub(channel, timeout=timeout, metadata=metadata)
+        )
+
+
+@pytest.mark.asyncio
+async def test_service_call_lower_level_with_overrides():
+    ITERATIONS = 99
+
+    # Setting deadline
+    deadline = grpclib.metadata.Deadline.from_timeout(22)
+    metadata = {"authorization": "12345"}
+    kwarg_deadline = grpclib.metadata.Deadline.from_timeout(28)
+    kwarg_metadata = {"authorization": "12345"}
+    async with ChannelFor(
+        [ExampleService(test_hook=_get_server_side_test(deadline, metadata))]
+    ) as channel:
+        stub = ExampleServiceStub(channel, deadline=deadline, metadata=metadata)
+        response = await stub._unary_unary(
+            "/service.ExampleService/DoThing", DoThingRequest(ITERATIONS), DoThingResponse,
+            deadline=kwarg_deadline,
+            metadata=kwarg_metadata,
+        )
+        assert response.successful_iterations == ITERATIONS
+
+    # Setting timeout
+    timeout = 99
+    deadline = grpclib.metadata.Deadline.from_timeout(timeout)
+    metadata = {"authorization": "12345"}
+    kwarg_timeout = 9000
+    kwarg_deadline = grpclib.metadata.Deadline.from_timeout(kwarg_timeout)
+    kwarg_metadata = {"authorization": "09876"}
+    async with ChannelFor(
+        [
+            ExampleService(
+                test_hook=_get_server_side_test(kwarg_deadline, kwarg_metadata)
+            )
+        ]
+    ) as channel:
+        stub = ExampleServiceStub(channel, deadline=deadline, metadata=metadata)
+        response = await stub._unary_unary(
+            "/service.ExampleService/DoThing", DoThingRequest(ITERATIONS), DoThingResponse,
+            timeout=kwarg_timeout,
+            metadata=kwarg_metadata,
+        )
+        assert response.successful_iterations == ITERATIONS