Add v2 stream tests

Author: blast-hardcheese

src/replit_river/v2/session.py

@@ -464,6 +464,7 @@ def commit(msg: TransportMessage) -> None:
             # Wake up backpressured writer
             backpressure_waiter, _ = self._streams[pending.streamId]
             backpressure_waiter.set()
+
         def get_next_pending() -> TransportMessage | None:
             if self._send_buffer:
                 return self._send_buffer[0]
@@ -923,8 +924,7 @@ async def _encode_stream() -> None:
                             yield error_deserializer(result["payload"])
                         except Exception:
                             logger.exception(
-                                "Error during stream "
-                                f"error deserialization: {result}"
+                                f"Error during stream error deserialization: {result}"
                             )
                         continue
                     yield response_deserializer(result["payload"])
@@ -1228,10 +1228,14 @@ async def _serve(
             logger.debug(f"_serve loop count={idx}")
             idx += 1
             ws = None
-            while (
-                state := get_state()
-            ) in ConnectingStates or (ws := get_ws()) is None:
-                logger.debug("_handle_messages_from_ws spinning while connecting, %r %r", ws, state)
+            while (state := get_state()) in ConnectingStates or (
+                ws := get_ws()
+            ) is None:
+                logger.debug(
+                    "_handle_messages_from_ws spinning while connecting, %r %r",
+                    ws,
+                    state,
+                )
                 await block_until_connected()
                 if state in TerminalStates:
                     break

tests/conftest.py

@@ -17,6 +17,7 @@
 pytest_plugins = [
     "tests.v1.river_fixtures.logging",
     "tests.v1.river_fixtures.clientserver",
+    "tests.v2.fixtures",
 ]
 
 HandlerKind = Literal["rpc", "subscription-stream", "upload-stream", "stream"]

tests/v2/codegen/snapshot/snapshots/test_basic_stream/__init__.py

@@ -0,0 +1,13 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from pydantic import BaseModel
+from typing import Literal
+
+import replit_river as river
+
+
+from .test_service import Test_ServiceService
+
+
+class StreamClient:
+    def __init__(self, client: river.v2.Client[Literal[None]]):
+        self.test_service = Test_ServiceService(client)

tests/v2/codegen/snapshot/snapshots/test_basic_stream/test_service/__init__.py

@@ -0,0 +1,44 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+from typing import Any
+import datetime
+
+from pydantic import TypeAdapter
+
+from replit_river.error_schema import RiverError, RiverErrorTypeAdapter
+import replit_river as river
+
+
+from .stream_method import (
+    Stream_MethodInit,
+    Stream_MethodInput,
+    Stream_MethodOutput,
+    Stream_MethodOutputTypeAdapter,
+    encode_Stream_MethodInit,
+    encode_Stream_MethodInput,
+)
+
+
+class Test_ServiceService:
+    def __init__(self, client: river.v2.Client[Any]):
+        self.client = client
+
+    async def stream_method(
+        self,
+        init: Stream_MethodInit,
+        inputStream: AsyncIterable[Stream_MethodInput],
+    ) -> AsyncIterator[Stream_MethodOutput | RiverError | RiverError]:
+        return self.client.send_stream(
+            "test_service",
+            "stream_method",
+            init,
+            inputStream,
+            encode_Stream_MethodInit,
+            encode_Stream_MethodInput,
+            lambda x: Stream_MethodOutputTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            lambda x: RiverErrorTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+        )

tests/v2/codegen/snapshot/snapshots/test_basic_stream/test_service/emit_error.py

@@ -0,0 +1,40 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+import datetime
+from typing import (
+    Any,
+    Literal,
+    Mapping,
+    NotRequired,
+    TypedDict,
+)
+from typing_extensions import Annotated
+
+from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
+from replit_river.error_schema import RiverError
+from replit_river.client import (
+    RiverUnknownError,
+    translate_unknown_error,
+    RiverUnknownValue,
+    translate_unknown_value,
+)
+
+import replit_river as river
+
+
+class Emit_ErrorErrorsOneOf_DATA_LOSS(RiverError):
+    code: Literal["DATA_LOSS"]
+    message: str
+
+
+class Emit_ErrorErrorsOneOf_UNEXPECTED_DISCONNECT(RiverError):
+    code: Literal["UNEXPECTED_DISCONNECT"]
+    message: str
+
+
+Emit_ErrorErrors = Annotated[
+    Emit_ErrorErrorsOneOf_DATA_LOSS
+    | Emit_ErrorErrorsOneOf_UNEXPECTED_DISCONNECT
+    | RiverUnknownError,
+    WrapValidator(translate_unknown_error),
+]

tests/v2/codegen/snapshot/snapshots/test_basic_stream/test_service/stream_method.py

@@ -0,0 +1,59 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+import datetime
+from typing import (
+    Any,
+    Literal,
+    Mapping,
+    NotRequired,
+    TypedDict,
+)
+from typing_extensions import Annotated
+
+from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
+from replit_river.error_schema import RiverError
+from replit_river.client import (
+    RiverUnknownError,
+    translate_unknown_error,
+    RiverUnknownValue,
+    translate_unknown_value,
+)
+
+import replit_river as river
+
+
+def encode_Stream_MethodInit(
+    _: "Stream_MethodInit",
+) -> Any:
+    return {}
+
+
+class Stream_MethodInit(TypedDict):
+    pass
+
+
+def encode_Stream_MethodInput(
+    x: "Stream_MethodInput",
+) -> Any:
+    return {
+        k: v
+        for (k, v) in (
+            {
+                "data": x.get("data"),
+            }
+        ).items()
+        if v is not None
+    }
+
+
+class Stream_MethodInput(TypedDict):
+    data: str
+
+
+class Stream_MethodOutput(BaseModel):
+    data: str
+
+
+Stream_MethodOutputTypeAdapter: TypeAdapter[Stream_MethodOutput] = TypeAdapter(
+    Stream_MethodOutput
+)

tests/v2/datagrams.py

@@ -0,0 +1,94 @@
+from dataclasses import dataclass
+from typing import (
+    Any,
+    NewType,
+    TypeAlias,
+)
+
+Datagram = dict[str, Any]
+TestTransport: TypeAlias = "FromClient | ToClient | WaitForClosed"
+
+StreamId = NewType("StreamId", str)
+ClientId = NewType("ClientId", str)
+ServerId = NewType("ServerId", str)
+SessionId = NewType("SessionId", str)
+
+
+@dataclass(frozen=True)
+class StreamAlias:
+    alias_id: int
+
+
+@dataclass(frozen=True)
+class ValueSet:
+    seq: int
+    ack: int
+    from_: ServerId | None = None
+    to: ClientId | None = None
+    procedureName: str | None = None
+    serviceName: str | None = None
+    create_alias: StreamAlias | None = None
+    stream_alias: StreamAlias | None = None
+    payload: Datagram | None = None
+
+
+@dataclass(frozen=True)
+class FromClient:
+    handshake_request: tuple[ClientId, ServerId, SessionId] | ValueSet | None = None
+    stream_open: tuple[ClientId, ServerId, str, str, StreamId] | ValueSet | None = None
+    stream_frame: tuple[ClientId, ServerId, int, int, Datagram] | ValueSet | None = None
+
+
+@dataclass(frozen=True)
+class ToClient:
+    seq: int
+    ack: int
+    control_flags: int = 0
+    handshake_response: bool | None = None
+    stream_frame: tuple[StreamAlias, Datagram] | None = None
+    stream_close: StreamAlias | None = None
+
+
+@dataclass(frozen=True)
+class WaitForClosed:
+    pass
+
+
+def decode_FromClient(datagram: dict[str, Any]) -> FromClient:
+    assert "from" in datagram
+    assert "to" in datagram
+    if datagram.get("payload", {}).get("type") == "HANDSHAKE_REQ":
+        assert "payload" in datagram
+        assert "sessionId" in datagram["payload"]
+        return FromClient(
+            handshake_request=(
+                ClientId(datagram["from"]),
+                ServerId(datagram["to"]),
+                SessionId(datagram["payload"]["sessionId"]),
+            )
+        )
+    elif datagram.get("controlFlags", 0) & 0b00010 > 0:  # STREAM_OPEN_BIT
+        return FromClient(
+            stream_open=(
+                ClientId(datagram["from"]),
+                ServerId(datagram["to"]),
+                datagram["serviceName"],
+                datagram["procedureName"],
+                StreamId(datagram["streamId"]),
+            )
+        )
+    elif datagram:
+        return FromClient(
+            stream_frame=(
+                ClientId(datagram["from"]),
+                ServerId(datagram["to"]),
+                datagram["seq"],
+                datagram["ack"],
+                datagram["payload"],
+            )
+        )
+    raise ValueError("Unexpected datagram: %r", datagram)
+
+
+def parser(datagram: dict[str, Any]) -> FromClient:
+    return decode_FromClient(datagram)