Add the typed client to the test harness

Author: igalshilman

examples/client.py

@@ -18,7 +18,6 @@
 
 
 async def main():
-    
     async with restate.create_client("http://localhost:8080") as client:
         await client.object_call(increment, key="a", arg=5)
 

examples/example.py

@@ -12,6 +12,11 @@
 # pylint: disable=C0116
 # pylint: disable=W0613
 
+
+#
+# uv run examples/example.py
+#
+
 import logging
 import restate
 

examples/harness.py

@@ -0,0 +1,38 @@
+#
+#  Copyright (c) 2023-2024 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+import restate
+
+from virtual_object import increment, count, counter
+
+#
+# uv run examples/harness.py
+#
+
+
+async def main():
+    app = restate.app([counter])
+    async with restate.create_test_harness(app) as harness:
+        await harness.client.object_call(increment, key="a", arg=5)
+        await harness.client.object_call(increment, key="a", arg=5)
+
+        current_count = await harness.client.object_call(count, key="a", arg=None)
+
+        print(f"Current count for 'a': {current_count}")
+
+        send = await harness.client.object_send(increment, key="b", arg=-10, idempotency_key="op1")
+
+        print(f"Sent increment to 'b', invocation ID: {send.invocation_id}")
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

examples/requirements.txt

@@ -25,7 +25,7 @@ test = ["pytest", "hypercorn"]
 lint = ["mypy>=1.11.2", "pyright>=1.1.390", "ruff>=0.6.9"]
 harness = ["testcontainers", "hypercorn", "httpx"]
 serde = ["dacite", "pydantic"]
-client = ["httpx"]
+client = ["httpx[http2]"]
 
 [build-system]
 requires = ["maturin>=1.6,<2.0"]

pyproject.toml

@@ -12,7 +12,7 @@
 Restate SDK for Python
 """
 
-from contextlib import contextmanager, asynccontextmanager
+from contextlib import asynccontextmanager
 import typing
 
 from restate.server_types import RestateAppT
@@ -51,14 +51,14 @@
     # we don't have the appropriate dependencies installed
 
     # pylint: disable=unused-argument, redefined-outer-name
-    @contextmanager
+    @asynccontextmanager
     def create_test_harness(
         app: RestateAppT,
         follow_logs: bool = False,
         restate_image: str = "restatedev/restate:latest",
         always_replay: bool = False,
         disable_retries: bool = False,
-    ) -> typing.Generator[TestHarnessEnvironment, None, None]:
+    ) -> typing.AsyncGenerator[TestHarnessEnvironment, None]:
         """a dummy harness constructor to raise ImportError. Install restate-sdk[harness] to use this feature"""
         raise ImportError("Install restate-sdk[harness] to use this feature")
 

python/restate/__init__.py

@@ -312,5 +312,5 @@ async def create_client(ingress: str, headers: typing.Optional[dict] = None) ->
     """
     Create a new Restate client.
     """
-    async with httpx.AsyncClient(base_url=ingress, headers=headers) as http_client:
+    async with httpx.AsyncClient(base_url=ingress, headers=headers, http2=True) as http_client:
         yield Client(http_client, headers)

python/restate/client.py

@@ -17,15 +17,17 @@
 import typing
 from urllib.error import URLError
 import socket
-from contextlib import contextmanager
+from contextlib import contextmanager, asynccontextmanager
 from warnings import deprecated
 
 from hypercorn.config import Config
 from hypercorn.asyncio import serve
+from restate.client import create_client
 from restate.server_types import RestateAppT
 from restate.types import TestHarnessEnvironment
 from testcontainers.core.container import DockerContainer  # type: ignore
 from testcontainers.core.waiting_utils import wait_container_is_ready  # type: ignore
+
 import httpx
 
 
@@ -332,14 +334,14 @@ def test_harness(
     return RestateTestHarness(app, config)
 
 
-@contextmanager
-def create_test_harness(
+@asynccontextmanager
+async def create_test_harness(
     app: RestateAppT,
     follow_logs: bool = False,
     restate_image: str = "restatedev/restate:latest",
     always_replay: bool = False,
     disable_retries: bool = False,
-) -> typing.Generator[TestHarnessEnvironment, None, None]:
+) -> typing.AsyncGenerator[TestHarnessEnvironment, None]:
     """
     Creates a test harness for running Restate services together with restate-server.
 
@@ -378,4 +380,7 @@ def create_test_harness(
             msg = f"unable to register the services at {bind_address} - {res.status_code} {res.text}"
             raise AssertionError(msg)
 
-        yield TestHarnessEnvironment(ingress_url=runtime.ingress_url(), admin_api_url=runtime.admin_url())
+        async with create_client(runtime.ingress_url()) as client:
+            yield TestHarnessEnvironment(
+                ingress_url=runtime.ingress_url(), admin_api_url=runtime.admin_url(), client=client
+            )

python/restate/harness.py

@@ -17,6 +17,8 @@
 
 from dataclasses import dataclass
 
+from restate.client import RestateClient
+
 
 @dataclass
 class TestHarnessEnvironment:
@@ -27,3 +29,6 @@ class TestHarnessEnvironment:
 
     admin_api_url: str
     """The URL of the Restate admin API endpoint used in the test"""
+
+    client: RestateClient
+    """The Restate client connected to the ingress URL"""