1212# pylint: disable=C0116
1313# pylint: disable=W0613
1414
15+ import os
1516from datetime import timedelta
16- from typing import Dict , Any
17- import typing
17+ from typing import (Dict , Iterable , List , Union , TypedDict , Literal , Any )
1818from restate import Service , Context
1919
20- from . import awakable_holder
20+ from . import list_object
21+ from . import awakeable_holder
2122
2223test_utils = Service ("TestUtilsService" )
2324
@@ -37,7 +38,7 @@ async def echo_headers(context: Context) -> Dict[str, str]:
3738async def create_awakeable_and_await_it (context : Context , req : Dict [str , Any ]) -> Dict [str , Any ]:
3839 name , awakeable = context .awakeable ()
3940
40- await context .object_call (awakable_holder .hold , key = req ["awakeableKey" ], arg = name )
41+ await context .object_call (awakeable_holder .hold , key = req ["awakeableKey" ], arg = name )
4142
4243 if "awaitTimeout" not in req :
4344 return {"type" : "result" , "value" : await awakeable }
@@ -46,7 +47,7 @@ async def create_awakeable_and_await_it(context: Context, req: Dict[str, Any]) -
4647 raise NotImplementedError ()
4748
4849@test_utils .handler (name = "sleepConcurrently" )
49- async def sleep_concurrently (context : Context , millis_duration : typing . List [int ]) -> None :
50+ async def sleep_concurrently (context : Context , millis_duration : List [int ]) -> None :
5051 timers = [context .sleep (timedelta (milliseconds = duration )) for duration in millis_duration ]
5152
5253 for timer in timers :
@@ -65,4 +66,35 @@ def effect():
6566 await context .run ("count" , effect )
6667
6768 return invoked_side_effects
68-
69+
70+ @test_utils .handler (name = "getEnvVariable" )
71+ async def get_env_variable (context : Context , env_name : str ) -> str :
72+ return os .environ .get (env_name , default = "" )
73+
74+ class CreateAwakeableAndAwaitIt (TypedDict ):
75+ type : Literal ["createAwakeableAndAwaitIt" ]
76+ awakeableKey : str
77+
78+ class GetEnvVariable (TypedDict ):
79+ type : Literal ["getEnvVariable" ]
80+ envName : str
81+
82+ Command = Union [
83+ CreateAwakeableAndAwaitIt ,
84+ GetEnvVariable
85+ ]
86+
87+ class InterpretRequest (TypedDict ):
88+ listName : str
89+ commands : Iterable [Command ]
90+
91+ @test_utils .handler (name = "interpretCommands" )
92+ async def interpret_commands (context : Context , req : InterpretRequest ):
93+ for cmd in req ['commands' ]:
94+ if cmd ['type' ] == "createAwakeableAndAwaitIt" :
95+ name , awakeable = context .awakeable ()
96+ context .object_send (awakeable_holder .hold , key = cmd ["awakeableKey" ], arg = name )
97+ result = await awakeable
98+ context .object_send (list_object .append , key = req ['listName' ], arg = result )
99+ elif cmd ['type' ] == "getEnvVariable" :
100+ context .object_send (list_object .append , key = req ['listName' ], arg = os .environ .get (cmd ['envName' ], default = "" ))
0 commit comments