Add some test

Author: cbrewster

replit_river/client.py

@@ -4,6 +4,7 @@
 from typing import Any, Generator, Generic, Literal, Optional, Union
 
 from opentelemetry import trace
+from opentelemetry.trace import Span, SpanKind, StatusCode
 
 from replit_river.client_transport import ClientTransport
 from replit_river.error_schema import RiverError, RiverException
@@ -118,7 +119,7 @@ async def send_subscription(
                 span,
             ):
                 if isinstance(msg, RiverError):
-                    _record_river_error(msg)
+                    _record_river_error(span, msg)
                 yield msg  # type: ignore # https://github.com/python/mypy/issues/10817
 
     async def send_stream(
@@ -146,30 +147,29 @@ async def send_stream(
                 span,
             ):
                 if isinstance(msg, RiverError):
-                    _record_river_error(msg)
+                    _record_river_error(span, msg)
                 yield msg  # type: ignore # https://github.com/python/mypy/issues/10817
 
 
-def _record_river_error(error: RiverError) -> None:
-    span = trace.get_current_span()
-    span.record_exception(RiverException(error.code, error.message))
-    span.set_attribute("river.error_code", error.code)
-    span.set_attribute("river.error_message", error.message)
-
-
 @contextmanager
 def _trace_procedure(
     procedure_type: Literal["rpc", "upload", "subscription", "stream"],
     service_name: str,
     procedure_name: str,
-) -> Generator[trace.Span, None, None]:
+) -> Generator[Span, None, None]:
     with tracer.start_span(
         f"river.client.{procedure_type}.{service_name}.{procedure_name}",
-        kind=trace.SpanKind.CLIENT,
+        kind=SpanKind.CLIENT,
     ) as span:
         try:
             yield span
         except RiverException as e:
-            span.set_attribute("river.error_code", e.code)
-            span.set_attribute("river.error_message", e.message)
+            _record_river_error(span, RiverError(code=e.code, message=e.message))
             raise e
+
+
+def _record_river_error(span: Span, error: RiverError) -> None:
+    span.set_status(StatusCode.ERROR, error.message)
+    span.record_exception(RiverException(error.code, error.message))
+    span.set_attribute("river.error_code", error.code)
+    span.set_attribute("river.error_message", error.message)

replit_river/rpc.py

@@ -388,8 +388,9 @@ async def _convert_outputs() -> None:
 
             convert_inputs_task = task_manager.create_task(_convert_inputs())
             convert_outputs_task = task_manager.create_task(_convert_outputs())
-            await asyncio.wait((convert_inputs_task, convert_outputs_task))
-
+            done, _ = await asyncio.wait((convert_inputs_task, convert_outputs_task))
+            for task in done:
+                await task
         except Exception as e:
             logger.exception("Uncaught exception in upload")
             await output.put(
@@ -440,17 +441,16 @@ async def _convert_inputs() -> None:
             response = method(request, context)
 
             async def _convert_outputs() -> None:
-                try:
-                    async for item in response:
-                        await output.put(
-                            get_response_or_error_payload(item, response_serializer)
-                        )
-                finally:
-                    output.close()
+                async for item in response:
+                    await output.put(
+                        get_response_or_error_payload(item, response_serializer)
+                    )
 
             convert_inputs_task = task_manager.create_task(_convert_inputs())
             convert_outputs_task = task_manager.create_task(_convert_outputs())
-            await asyncio.wait((convert_inputs_task, convert_outputs_task))
+            done, _ = await asyncio.wait((convert_inputs_task, convert_outputs_task))
+            for task in done:
+                await task
         except grpc.RpcError:
             logger.exception("RPC exception in stream")
             code = grpc.StatusCode(context._abort_code).name if context else "UNKNOWN"

tests/conftest.py

@@ -3,15 +3,19 @@
 from collections.abc import AsyncIterator
 from typing import Any, AsyncGenerator, Iterator, Literal
 
+import grpc.aio
 import nanoid  # type: ignore
 import pytest
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from websockets.server import serve
 
 from replit_river.client import Client
 from replit_river.client_transport import UriAndMetadata
-from replit_river.error_schema import RiverError
+from replit_river.error_schema import RiverError, RiverException
 from replit_river.rpc import (
-    GrpcContext,
     TransportMessage,
     rpc_method_handler,
     stream_method_handler,
@@ -68,12 +72,12 @@ def deserialize_error(response: dict) -> RiverError:
 
 
 # RPC method handlers for testing
-async def rpc_handler(request: str, context: GrpcContext) -> str:
+async def rpc_handler(request: str, context: grpc.aio.ServicerContext) -> str:
     return f"Hello, {request}!"
 
 
 async def subscription_handler(
-    request: str, context: GrpcContext
+    request: str, context: grpc.aio.ServicerContext
 ) -> AsyncGenerator[str, None]:
     for i in range(5):
         yield f"Subscription message {i} for {request}"
@@ -93,7 +97,8 @@ async def upload_handler(
 
 
 async def stream_handler(
-    request: Iterator[str] | AsyncIterator[str], context: GrpcContext
+    request: Iterator[str] | AsyncIterator[str],
+    context: grpc.aio.ServicerContext,
 ) -> AsyncGenerator[str, None]:
     if isinstance(request, AsyncIterator):
         async for data in request:
@@ -103,6 +108,14 @@ async def stream_handler(
             yield f"Stream response for {data}"
 
 
+async def stream_error_handler(
+    request: Iterator[str] | AsyncIterator[str],
+    context: grpc.aio.ServicerContext,
+) -> AsyncGenerator[str, None]:
+    raise RiverException("INJECTED_ERROR", "test error")
+    yield "test"  # appease the type checker
+
+
 @pytest.fixture
 def transport_options() -> TransportOptions:
     return TransportOptions()
@@ -137,6 +150,12 @@ def server(transport_options: TransportOptions) -> Server:
                     stream_handler, deserialize_request, serialize_response
                 ),
             ),
+            ("test_service", "stream_method_error"): (
+                "stream",
+                stream_method_handler(
+                    stream_error_handler, deserialize_request, serialize_response
+                ),
+            ),
         }
     )
     return server
@@ -173,3 +192,18 @@ async def websocket_uri_factory() -> UriAndMetadata[None]:
         await server.close()
         # Server should close normally
         no_logging_error()
+
+
+@pytest.fixture(scope="session")
+def span_exporter() -> InMemorySpanExporter:
+    exporter = InMemorySpanExporter()
+    processor = SimpleSpanProcessor(exporter)
+    provider = TracerProvider()
+    provider.add_span_processor(processor)
+    trace.set_tracer_provider(provider)
+    return exporter
+
+
+@pytest.fixture(autouse=True)
+def reset_span_exporter(span_exporter: InMemorySpanExporter) -> None:
+    span_exporter.clear()

tests/river_fixtures/logging.py

@@ -15,8 +15,15 @@ class NoErrors:
 
     def __init__(self, caplog: LogCaptureFixture):
         self.caplog = caplog
+        self._allow_errors = False
+
+    def allow_errors(self) -> None:
+        self._allow_errors = True
 
     def __call__(self) -> None:
+        if self._allow_errors:
+            return
+
         assert len(self.caplog.get_records("setup")) == 0
         assert len(self.caplog.get_records("call")) == 0
         assert len(self.caplog.get_records("teardown")) == 0

tests/test_communication.py

@@ -18,7 +18,7 @@ async def test_rpc_method(client: Client) -> None:
         serialize_request,
         deserialize_response,
         deserialize_error,
-    )  # type: ignore
+    )
     assert response == "Hello, Alice!"
 
 

tests/test_opentelemetry.py

@@ -0,0 +1,143 @@
+from typing import AsyncGenerator
+
+import pytest
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import StatusCode
+
+from replit_river.client import Client
+from replit_river.error_schema import RiverError
+from tests.conftest import deserialize_error, deserialize_response, serialize_request
+from tests.river_fixtures.logging import NoErrors
+
+
+@pytest.mark.asyncio
+async def test_rpc_method_span(
+    client: Client, span_exporter: InMemorySpanExporter
+) -> None:
+    response = await client.send_rpc(
+        "test_service",
+        "rpc_method",
+        "Alice",
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+    )
+    assert response == "Hello, Alice!"
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "river.client.rpc.test_service.rpc_method"
+
+
+@pytest.mark.asyncio
+async def test_upload_method_span(
+    client: Client, span_exporter: InMemorySpanExporter
+) -> None:
+    async def upload_data() -> AsyncGenerator[str, None]:
+        yield "Data 1"
+        yield "Data 2"
+        yield "Data 3"
+
+    response = await client.send_upload(
+        "test_service",
+        "upload_method",
+        "Initial Data",
+        upload_data(),
+        serialize_request,
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+    )
+    assert response == "Uploaded: Initial Data, Data 1, Data 2, Data 3"
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "river.client.upload.test_service.upload_method"
+
+
+@pytest.mark.asyncio
+async def test_subscription_method_span(
+    client: Client, span_exporter: InMemorySpanExporter
+) -> None:
+    async for response in client.send_subscription(
+        "test_service",
+        "subscription_method",
+        "Bob",
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+    ):
+        assert isinstance(response, str)
+        assert "Subscription message" in response
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "river.client.subscription.test_service.subscription_method"
+
+
+@pytest.mark.asyncio
+async def test_stream_method_span(
+    client: Client, span_exporter: InMemorySpanExporter
+) -> None:
+    async def stream_data() -> AsyncGenerator[str, None]:
+        yield "Stream 1"
+        yield "Stream 2"
+        yield "Stream 3"
+
+    responses = []
+    async for response in client.send_stream(
+        "test_service",
+        "stream_method",
+        "Initial Stream Data",
+        stream_data(),
+        serialize_request,
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+    ):
+        responses.append(response)
+
+    assert responses == [
+        "Stream response for Initial Stream Data",
+        "Stream response for Stream 1",
+        "Stream response for Stream 2",
+        "Stream response for Stream 3",
+    ]
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "river.client.stream.test_service.stream_method"
+
+
+@pytest.mark.asyncio
+async def test_stream_error_method_span(
+    client: Client,
+    span_exporter: InMemorySpanExporter,
+    no_logging_error: NoErrors,
+) -> None:
+    # We are explicitly testing errors.
+    no_logging_error.allow_errors()
+
+    async def stream_data() -> AsyncGenerator[str, None]:
+        yield "Stream 1"
+        yield "Stream 2"
+        yield "Stream 3"
+
+    responses = []
+    async for response in client.send_stream(
+        "test_service",
+        "stream_method_error",
+        "Initial Stream Data",
+        stream_data(),
+        serialize_request,
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+    ):
+        responses.append(response)
+
+    assert len(responses) == 1
+    assert isinstance(responses[0], RiverError)
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "river.client.stream.test_service.stream_method_error"
+    assert spans[0].status.status_code == StatusCode.ERROR