[WIP] update test suite to 3.0

Author: slinkydeveloper

.github/workflows/integration.yaml

@@ -104,7 +104,7 @@ jobs:
           cache-to: type=gha,mode=max,scope=${{ github.workflow }}
 
       - name: Run test tool
-        uses: restatedev/sdk-test-suite@v2.4
+        uses: restatedev/sdk-test-suite@v3.0
         with:
           restateContainerImage: ${{ inputs.restateCommit != '' && 'localhost/restatedev/restate-commit-download:latest' || (inputs.restateImage != '' && inputs.restateImage || 'ghcr.io/restatedev/restate:main') }}
           serviceContainerImage: "restatedev/python-test-services"

test-services/exclusions.yaml

@@ -1,15 +1 @@
-exclusions:
-  "alwaysSuspending":
-  - "dev.restate.sdktesting.tests.AwaitTimeout"
-  "default":
-  - "dev.restate.sdktesting.tests.AwaitTimeout"
-  - "dev.restate.sdktesting.tests.RawHandler"
-  "lazyState": []
-  "singleThreadSinglePartition":
-  - "dev.restate.sdktesting.tests.AwaitTimeout"
-  - "dev.restate.sdktesting.tests.RawHandler"
-  "threeNodes":
-  - "dev.restate.sdktesting.tests.AwaitTimeout"
-  - "dev.restate.sdktesting.tests.RawHandler"
-  "threeNodesAlwaysSuspending":
-  - "dev.restate.sdktesting.tests.AwaitTimeout"
+exclusions: {}

test-services/services/__init__.py

@@ -22,6 +22,7 @@
 from .map_object import map_object as s9
 from .non_determinism import non_deterministic as s10
 from .test_utils import test_utils as s11
+from .virtual_object_command_interpreter import virtual_object_command_interpreter as s16
 
 from .interpreter import layer_0 as s12
 from .interpreter import layer_1 as s13

test-services/services/cancel_test.py

@@ -27,7 +27,7 @@
 @runner.handler(name="startTest")
 async def start_test(ctx: ObjectContext, op: BlockingOperation):
     try:
-        await ctx.object_call(block, key="", arg=op)
+        await ctx.object_call(block, key=ctx.key(), arg=op)
     except TerminalError as t:
         if t.status_code == 409:
             ctx.set("state", True)
@@ -47,11 +47,11 @@ async def verify_test(ctx: ObjectContext) -> bool:
 @blocking_service.handler()
 async def block(ctx: ObjectContext, op: BlockingOperation):
     name, awakeable = ctx.awakeable()
-    await ctx.object_call(awakeable_holder.hold, key="cancel", arg=name)
+    await ctx.object_call(awakeable_holder.hold, key=ctx.key(), arg=name)
     await awakeable
 
     if op == "CALL":
-        await ctx.object_call(block, key="", arg=op)
+        await ctx.object_call(block, key=ctx.key(), arg=op)
     elif op == "SLEEP":
         await ctx.sleep(timedelta(days=1024))
     elif op == "AWAKEABLE":

test-services/services/kill_test.py

@@ -12,25 +12,25 @@
 # pylint: disable=C0116
 # pylint: disable=W0613
 
-from restate import Service, Context, VirtualObject, ObjectContext
+from restate import VirtualObject, ObjectContext
 
 from . import awakeable_holder
 
-kill_runner = Service("KillTestRunner")
+kill_runner = VirtualObject("KillTestRunner")
 
 @kill_runner.handler(name="startCallTree")
-async def start_call_tree(ctx: Context):
-    await ctx.object_call(recursive_call, key="", arg=None)
+async def start_call_tree(ctx: ObjectContext):
+    await ctx.object_call(recursive_call, key=ctx.key(), arg=None)
 
 kill_singleton = VirtualObject("KillTestSingleton")
 
 @kill_singleton.handler(name="recursiveCall")
 async def recursive_call(ctx: ObjectContext):
     name, promise = ctx.awakeable()
-    ctx.object_send(awakeable_holder.hold, key="kill", arg=name)
+    ctx.object_send(awakeable_holder.hold, key=ctx.key(), arg=name)
     await promise
 
-    await ctx.object_call(recursive_call, key="", arg=None)
+    await ctx.object_call(recursive_call, key=ctx.key(), arg=None)
 
 @kill_singleton.handler(name="isUnlocked")
 async def is_unlocked(ctx: ObjectContext):

test-services/services/proxy.py

@@ -25,6 +25,7 @@ class ProxyRequest(TypedDict):
     handlerName: str
     message: Iterable[int]
     delayMillis: Optional[int]
+    idempotencyKey: Optional[str]
 
 
 @proxy.handler()
@@ -33,7 +34,8 @@ async def call(ctx: Context, req: ProxyRequest) -> Iterable[int]:
         req['serviceName'],
         req['handlerName'],
         bytes(req['message']),
-        req.get('virtualObjectKey')))
+        req.get('virtualObjectKey'),
+        req.get('idempotencyKey')))
 
 
 @proxy.handler(name="oneWayCall")
@@ -46,7 +48,8 @@ async def one_way_call(ctx: Context, req: ProxyRequest):
         req['handlerName'],
         bytes(req['message']),
         req.get('virtualObjectKey'),
-        send_delay
+        send_delay,
+        req.get('idempotencyKey')
     )
 
 
@@ -69,14 +72,16 @@ async def many_calls(ctx: Context, requests: Iterable[ManyCallRequest]):
                 req['proxyRequest']['handlerName'],
                 bytes(req['proxyRequest']['message']),
                 req['proxyRequest'].get('virtualObjectKey'),
-                send_delay
+                send_delay,
+                req['proxyRequest'].get('idempotencyKey')
             )
         else:
             awaitable = ctx.generic_call(
                 req['proxyRequest']['serviceName'],
                 req['proxyRequest']['handlerName'],
                 bytes(req['proxyRequest']['message']),
-                req['proxyRequest'].get('virtualObjectKey'))
+                req['proxyRequest'].get('virtualObjectKey'),
+                req['proxyRequest'].get('idempotencyKey'))
             if req['awaitAtTheEnd']:
                 to_await.append(awaitable)
 

test-services/services/test_utils.py

@@ -34,18 +34,6 @@ async def uppercase_echo(context: Context, input: str) -> str:
 async def echo_headers(context: Context) -> Dict[str, str]:
     return context.request().headers
 
-@test_utils.handler(name="createAwakeableAndAwaitIt")
-async def create_awakeable_and_await_it(context: Context, req: Dict[str, Any]) -> Dict[str, Any]:
-    name, awakeable = context.awakeable()
-
-    await context.object_call(awakeable_holder.hold, key=req["awakeableKey"], arg=name)
-
-    if "awaitTimeout" not in req:
-        return {"type": "result", "value": await awakeable}
-
-    timeout = context.sleep(timedelta(milliseconds=int(req["awaitTimeout"])))
-    raise NotImplementedError()
-
 @test_utils.handler(name="sleepConcurrently")
 async def sleep_concurrently(context: Context, millis_duration: List[int]) -> None:
     timers = [context.sleep(timedelta(milliseconds=duration)) for duration in millis_duration]
@@ -67,34 +55,6 @@ def effect():
 
     return invoked_side_effects
 
-@test_utils.handler(name="getEnvVariable")
-async def get_env_variable(context: Context, env_name: str) -> str:
-    return os.environ.get(env_name, default="")
-
-class CreateAwakeableAndAwaitIt(TypedDict):
-    type: Literal["createAwakeableAndAwaitIt"]
-    awakeableKey: str
-
-class GetEnvVariable(TypedDict):
-    type: Literal["getEnvVariable"]
-    envName: str
-
-Command = Union[
-    CreateAwakeableAndAwaitIt,
-    GetEnvVariable
-]
-
-class InterpretRequest(TypedDict):
-    listName: str
-    commands: Iterable[Command]
-
-@test_utils.handler(name="interpretCommands")
-async def interpret_commands(context: Context, req: InterpretRequest):
-    for cmd in req['commands']:
-        if cmd['type'] == "createAwakeableAndAwaitIt":
-            name, awakeable = context.awakeable()
-            context.object_send(awakeable_holder.hold, key=cmd["awakeableKey"], arg=name)
-            result = await awakeable
-            context.object_send(list_object.append, key=req['listName'], arg=result)
-        elif cmd['type'] == "getEnvVariable":
-            context.object_send(list_object.append, key=req['listName'], arg=os.environ.get(cmd['envName'], default=""))
+@test_utils.handler(name="cancelInvocation")
+async def cancel_invocation(context: Context, invocation_id: str) -> None:
+    context.cancel(invocation_id)

test-services/services/virtual_object_command_interpreter.py

@@ -0,0 +1,138 @@
+#
+#  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
+#
+"""example.py"""
+# pylint: disable=C0116
+# pylint: disable=W0613
+
+import os
+from datetime import timedelta
+from typing import (Dict, Iterable, List, Union, TypedDict, Literal, Any)
+from restate import VirtualObject, ObjectSharedContext, ObjectContext
+from restate import select
+from restate.exceptions import TerminalError
+
+virtual_object_command_interpreter = VirtualObject("VirtualObjectCommandInterpreter")
+
+@virtual_object_command_interpreter.handler(name="getResults", kind="shared")
+async def get_results(ctx: ObjectSharedContext) -> List[str]:
+    return ctx.get("results") or []
+
+@virtual_object_command_interpreter.handler(name="hasAwakeable", kind="shared")
+async def has_awakeable(ctx: ObjectSharedContext, awk_key: str) -> bool:
+    awk_id = ctx.get("awk-" + awk_key)
+    if awk_id:
+        return True
+    return False
+
+class CreateAwakeable(TypedDict):
+    type: Literal["createAwakeable"]
+    awakeableKey: str
+
+class Sleep(TypedDict):
+    type: Literal["sleep"]
+    timeoutMillis: int
+
+class RunThrowTerminalException(TypedDict):
+    type: Literal["runThrowTerminalException"]
+    reason: str
+
+AwaitableCommand = Union[
+    CreateAwakeable,
+    Sleep,
+    RunThrowTerminalException
+]
+
+class AwaitOne(TypedDict):
+    type: Literal["awaitOne"]
+    command: AwaitableCommand
+
+class AwaitAnySuccessful(TypedDict):
+    type: Literal["awaitAnySuccessful"]
+    commands: List[AwaitableCommand]
+
+class AwaitAny(TypedDict):
+    type: Literal["awaitAny"]
+    commands: List[AwaitableCommand]
+
+class AwaitAwakeableOrTimeout(TypedDict):
+    type: Literal["awaitAwakeableOrTimeout"]
+    awakeableKey: str
+    timeoutMillis: int
+
+class ResolveAwakeable(TypedDict):
+    type: Literal["resolveAwakeable"]
+    awakeableKey: str
+    value: str
+
+class RejectAwakeable(TypedDict):
+    type: Literal["rejectAwakeable"]
+    awakeableKey: str
+    reason: str
+
+class GetEnvVariable(TypedDict):
+    type: Literal["getEnvVariable"]
+    envName: str
+
+Command = Union[
+    AwaitOne,
+    AwaitAny,
+    AwaitAnySuccessful,
+    AwaitAwakeableOrTimeout,
+    ResolveAwakeable,
+    RejectAwakeable,
+    GetEnvVariable
+]
+
+class InterpretRequest(TypedDict):
+    commands: Iterable[Command]
+
+@virtual_object_command_interpreter.handler(name="resolveAwakeable", kind="shared")
+async def resolve_awakeable(ctx: ObjectSharedContext, req: ResolveAwakeable) -> bool:
+    awk_id = ctx.get("awk-" + req.awakeableKey)
+    if not awk_id:
+        raise TerminalError(message="No awakeable is registered")
+    ctx.resolve_awakeable(awk_id, req.value)
+
+@virtual_object_command_interpreter.handler(name="rejectAwakeable", kind="shared")
+async def reject_awakeable(ctx: ObjectSharedContext, req: RejectAwakeable) -> bool:
+    awk_id = ctx.get("awk-" + req.awakeableKey)
+    if not awk_id:
+        raise TerminalError(message="No awakeable is registered")
+    ctx.reject_awakeable(awk_id, req.reason)
+
+@virtual_object_command_interpreter.handler(name="interpretCommands")
+async def interpret_commands(ctx: ObjectContext, req: InterpretRequest):
+    result = ""
+
+    for cmd in req['commands']:
+        if cmd['type'] == "awaitAwakeableOrTimeout":
+            awk_id, awakeable = ctx.awakeable()
+            ctx.get("awk-" + cmd.awakeableKey, awk_id)
+            match await select(awakeable=awakeable, timeout=ctx.sleep(timedelta(milliseconds=cmd.timeoutMillis))):
+                case ['awakeable', awk_res]:
+                    result = awk_res
+                case ['timeout', _]:
+                    raise TerminalError(message="await-timeout", status_code=500)
+        elif cmd['type'] == "resolveAwakeable":
+            resolve_awakeable(ctx, cmd)
+            result = ""
+        elif cmd['type'] == "rejectAwakeable":
+            reject_awakeable(ctx, cmd)
+            result = ""
+        elif cmd['type'] == "getEnvVariable":
+            result = await ctx.run("get_env", lambda: os.environ.get(cmd['envName'], default=""))
+
+        last_results = get_results(ctx)
+        last_results.append(result)
+        ctx.set("results", last_results)
+
+    return result
+