Fixes per PR review feedback.

Refactored workflow to use workflow.init.
Added unit test.

Author: mfateev

tests/updatable_timer/__init__.py

@@ -0,0 +1,33 @@
+import datetime
+import logging
+import math
+import uuid
+
+from temporalio.client import Client, WorkflowExecutionStatus
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from updatable_timer.workflow import Workflow
+
+
+async def test_updatable_timer_workflow(client: Client):
+    logging.basicConfig(level=logging.DEBUG)
+
+    task_queue_name = str(uuid.uuid4())
+    async with await WorkflowEnvironment.start_time_skipping() as env:
+        async with Worker(env.client, task_queue=task_queue_name, workflows=[Workflow]):
+            in_a_day = float(
+                (datetime.datetime.now() + datetime.timedelta(days=1)).timestamp()
+            )
+            in_an_hour = float(
+                (datetime.datetime.now() + datetime.timedelta(hours=1)).timestamp()
+            )
+            handle = await env.client.start_workflow(
+                Workflow.run, in_a_day, id=str(uuid.uuid4()), task_queue=task_queue_name
+            )
+            wake_up_time1 = await handle.query(Workflow.get_wake_up_time)
+            assert math.isclose(wake_up_time1, in_a_day)
+            await handle.signal(Workflow.update_wake_up_time, in_an_hour)
+            wake_up_time2 = await handle.query(Workflow.get_wake_up_time)
+            assert math.isclose(wake_up_time2, in_an_hour)
+            await handle.result()

tests/updatable_timer/updatable_timer_test.py

@@ -8,28 +8,25 @@
 
 @workflow.defn
 class Workflow:
-    def __init__(self):
-        self.timer: Optional[UpdatableTimer] = None
-
-    @workflow.run
-    async def run(self, wake_up_time: float):
+    @workflow.init
+    def __init__(self, wake_up_time: float) -> None:
         self.timer = UpdatableTimer(
             datetime.fromtimestamp(wake_up_time, tz=timezone.utc)
         )
+
+    @workflow.run
+    async def run(self, wake_up_time: float):
         await self.timer.sleep()
 
     @workflow.signal
-    async def update_wake_up_time(self, wake_up_time: float):
-        # Deals with situation when the signal method is called before the run method.
-        # This happens when a workflow task is executed after a signal is received
-        # or when a workflow is started using the signal-with-start.
-        await workflow.wait_condition(lambda: self.timer is not None)
-        assert self.timer is not None  # for mypy
+    async def update_wake_up_time(self, wake_up_time: float) -> None:
+        workflow.logger.info(f"update_wake_up_time: {wake_up_time}")
+
         self.timer.update_wake_up_time(
             datetime.fromtimestamp(wake_up_time, tz=timezone.utc)
         )
 
     @workflow.query
-    def get_wake_up_time(self):
-        assert self.timer is not None  # for mypy
-        return self.timer.get_wake_up_time()
+    def get_wake_up_time(self) -> float:
+        workflow.logger.info(f"get_wake_up_time")
+        return float(self.timer.get_wake_up_time().timestamp())