Add v2 rpc tests

Author: blast-hardcheese

tests/v2/codegen/snapshot/snapshots/test_basic_rpc/__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_rpc/test_service/__init__.py

@@ -0,0 +1,41 @@
+# 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 .rpc_method import (
+    Rpc_MethodInit,
+    Rpc_MethodOutput,
+    Rpc_MethodOutputTypeAdapter,
+    encode_Rpc_MethodInit,
+)
+
+
+class Test_ServiceService:
+    def __init__(self, client: river.v2.Client[Any]):
+        self.client = client
+
+    async def rpc_method(
+        self,
+        init: Rpc_MethodInit,
+        timeout: datetime.timedelta,
+    ) -> Rpc_MethodOutput:
+        return await self.client.send_rpc(
+            "test_service",
+            "rpc_method",
+            init,
+            encode_Rpc_MethodInit,
+            lambda x: Rpc_MethodOutputTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            lambda x: RiverErrorTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            timeout,
+        )

tests/v2/codegen/snapshot/snapshots/test_basic_rpc/test_service/rpc_method.py

@@ -0,0 +1,49 @@
+# 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_Rpc_MethodInit(
+    x: "Rpc_MethodInit",
+) -> Any:
+    return {
+        k: v
+        for (k, v) in (
+            {
+                "data": x.get("data"),
+            }
+        ).items()
+        if v is not None
+    }
+
+
+class Rpc_MethodInit(TypedDict):
+    data: str
+
+
+class Rpc_MethodOutput(BaseModel):
+    data: str
+
+
+Rpc_MethodOutputTypeAdapter: TypeAdapter[Rpc_MethodOutput] = TypeAdapter(
+    Rpc_MethodOutput
+)

tests/v2/test_rpc.py

@@ -0,0 +1,97 @@
+import importlib
+import logging
+from collections import deque
+from datetime import timedelta
+from typing import (
+    Literal,
+)
+
+import pytest
+from pytest_snapshot.plugin import Snapshot
+
+from replit_river.v2.client import Client
+from tests.fixtures.codegen_snapshot_fixtures import validate_codegen
+from tests.v2.datagrams import (
+    ClientId,
+    FromClient,
+    ServerId,
+    StreamAlias,
+    TestTransport,
+    ToClient,
+    ValueSet,
+    WaitForClosed,
+)
+
+logger = logging.getLogger(__name__)
+
+_AlreadyGenerated = False
+
+
+@pytest.fixture(scope="function", autouse=True)
+def stream_client_codegen(snapshot: Snapshot) -> Literal[True]:
+    global _AlreadyGenerated
+    if not _AlreadyGenerated:
+        validate_codegen(
+            snapshot=snapshot,
+            snapshot_dir="tests/v2/codegen/snapshot/snapshots",
+            read_schema=lambda: open("tests/v2/test_rpc.schema.json"),
+            target_path="test_basic_rpc",
+            client_name="StreamClient",
+            protocol_version="v2.0",
+        )
+        _AlreadyGenerated = True
+
+    import tests.v2.codegen.snapshot.snapshots.test_basic_stream
+
+    importlib.reload(tests.v2.codegen.snapshot.snapshots.test_basic_stream)
+    return True
+
+
+rpc_expected: deque[TestTransport] = deque(
+    [
+        FromClient(
+            handshake_request=ValueSet(
+                seq=0,  # These don't count due to being during a handshake
+                ack=0,
+                from_=ServerId("server-001"),
+                to=ClientId("client-001"),
+            )
+        ),
+        FromClient(
+            stream_open=ValueSet(
+                seq=0,
+                ack=0,
+                from_=ServerId("server-001"),
+                to=ClientId("client-001"),
+                serviceName="test_service",
+                procedureName="rpc_method",
+                create_alias=StreamAlias(1),
+                payload={"data": "foo"},
+                stream_closed=True,
+            )
+        ),
+        ToClient(
+            seq=0,
+            ack=1,
+            stream_frame=(
+                StreamAlias(1),
+                {"ok": True, "payload": {"data": "Hello, foo!"}},
+            ),
+        ),
+        WaitForClosed(),
+    ]
+)
+
+
+@pytest.mark.parametrize("expected", [rpc_expected])
+async def test_rpc(bound_client: Client) -> None:
+    from tests.v2.codegen.snapshot.snapshots.test_basic_rpc import (
+        StreamClient,
+    )
+
+    res = await StreamClient(bound_client).test_service.rpc_method(
+        init={"data": "foo"},
+        timeout=timedelta(seconds=5),
+    )
+
+    assert res.data == "Hello, foo!"

tests/v2/test_rpc.schema.json

@@ -0,0 +1,32 @@
+{
+  "services": {
+    "test_service": {
+      "procedures": {
+        "rpc_method": {
+          "init": {
+            "type": "object",
+            "properties": {
+              "data": {
+                "type": "string"
+              }
+            },
+            "required": ["data"]
+          },
+          "output": {
+            "type": "object",
+            "properties": {
+              "data": {
+                "type": "string"
+              }
+            },
+            "required": ["data"]
+          },
+          "errors": {
+            "not": {}
+          },
+          "type": "rpc"
+        }
+      }
+    }
+  }
+}