From b78db5354c77185b0698f125fd39d706b3587eb3 Mon Sep 17 00:00:00 2001
From: Devon Stewart <devon.stewart@repl.it>
Date: Thu, 28 Nov 2024 23:26:35 -0800
Subject: [PATCH] Adding an RPC codegen test

---
 tests/codegen/rpc/generated/__init__.py       | 13 ++++++
 .../rpc/generated/test_service/__init__.py    | 36 +++++++++++++++++
 .../rpc/generated/test_service/rpc_method.py  | 40 +++++++++++++++++++
 tests/codegen/rpc/schema.json                 | 32 +++++++++++++++
 tests/codegen/test_rpc.py                     | 24 +++++++++++
 5 files changed, 145 insertions(+)
 create mode 100644 tests/codegen/rpc/generated/__init__.py
 create mode 100644 tests/codegen/rpc/generated/test_service/__init__.py
 create mode 100644 tests/codegen/rpc/generated/test_service/rpc_method.py
 create mode 100644 tests/codegen/rpc/schema.json
 create mode 100644 tests/codegen/test_rpc.py

diff --git a/tests/codegen/rpc/generated/__init__.py b/tests/codegen/rpc/generated/__init__.py
new file mode 100644
index 0000000..efcb731
--- /dev/null
+++ b/tests/codegen/rpc/generated/__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 RpcClient:
+    def __init__(self, client: river.Client[Literal[None]]):
+        self.test_service = Test_ServiceService(client)
diff --git a/tests/codegen/rpc/generated/test_service/__init__.py b/tests/codegen/rpc/generated/test_service/__init__.py
new file mode 100644
index 0000000..da1bce5
--- /dev/null
+++ b/tests/codegen/rpc/generated/test_service/__init__.py
@@ -0,0 +1,36 @@
+# 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
+import replit_river as river
+
+
+from .rpc_method import encode_Rpc_MethodInput, Rpc_MethodInput, Rpc_MethodOutput
+
+
+class Test_ServiceService:
+    def __init__(self, client: river.Client[Any]):
+        self.client = client
+
+    async def rpc_method(
+        self,
+        input: Rpc_MethodInput,
+        timeout: datetime.timedelta,
+    ) -> Rpc_MethodOutput:
+        return await self.client.send_rpc(
+            'test_service',
+            'rpc_method',
+            input,
+            encode_Rpc_MethodInput,
+            lambda x: TypeAdapter(Rpc_MethodOutput).validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            lambda x: TypeAdapter(RiverError).validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            timeout,
+        )
diff --git a/tests/codegen/rpc/generated/test_service/rpc_method.py b/tests/codegen/rpc/generated/test_service/rpc_method.py
new file mode 100644
index 0000000..90677e5
--- /dev/null
+++ b/tests/codegen/rpc/generated/test_service/rpc_method.py
@@ -0,0 +1,40 @@
+# ruff: noqa
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+import datetime
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Mapping,
+    Union,
+    Tuple,
+    TypedDict,
+)
+
+from pydantic import BaseModel, Field, TypeAdapter
+from replit_river.error_schema import RiverError
+
+import replit_river as river
+
+
+encode_Rpc_MethodInput: Callable[['Rpc_MethodInput'], Any] = lambda x: {
+    k: v
+    for (k, v) in (
+        {
+            'data': x.get('data'),
+        }
+    ).items()
+    if v is not None
+}
+
+
+class Rpc_MethodInput(TypedDict):
+    data: str
+
+
+class Rpc_MethodOutput(BaseModel):
+    data: str
diff --git a/tests/codegen/rpc/schema.json b/tests/codegen/rpc/schema.json
new file mode 100644
index 0000000..508d354
--- /dev/null
+++ b/tests/codegen/rpc/schema.json
@@ -0,0 +1,32 @@
+{
+  "services": {
+    "test_service": {
+      "procedures": {
+        "rpc_method": {
+          "input": {
+            "type": "object",
+            "properties": {
+              "data": {
+                "type": "string"
+              }
+            },
+            "required": ["data"]
+          },
+          "output": {
+            "type": "object",
+            "properties": {
+              "data": {
+                "type": "string"
+              }
+            },
+            "required": ["data"]
+          },
+          "errors": {
+            "not": {}
+          },
+          "type": "rpc"
+        }
+      }
+    }
+  }
+}
diff --git a/tests/codegen/test_rpc.py b/tests/codegen/test_rpc.py
new file mode 100644
index 0000000..86f2028
--- /dev/null
+++ b/tests/codegen/test_rpc.py
@@ -0,0 +1,24 @@
+from datetime import timedelta
+
+import pytest
+
+from replit_river.client import Client
+from replit_river.codegen.client import schema_to_river_client_codegen
+from tests.common_handlers import basic_rpc_method
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("handlers", [{**basic_rpc_method}])
+async def test_basic_rpc(client: Client) -> None:
+    schema_to_river_client_codegen(
+        "tests/codegen/rpc/schema.json", "tests/codegen/rpc/generated", "RpcClient", True
+    )
+    from tests.codegen.rpc.generated import RpcClient
+
+    res = await RpcClient(client).test_service.rpc_method(
+        {
+            "data": "feep",
+        },
+        timedelta(seconds=5),
+    )
+    assert res.data == "Hello, feep!", f"Expected 'Hello, feep!' received {res.data}"