Merge branch 'main' into remove-update-experimental-notices

Author: dandavison

tests/helpers/__init__.py

@@ -4,7 +4,7 @@
 import uuid
 from contextlib import closing
 from datetime import timedelta
-from typing import Awaitable, Callable, Optional, Sequence, Type, TypeVar
+from typing import Any, Awaitable, Callable, Optional, Sequence, Type, TypeVar
 
 from temporalio.api.common.v1 import WorkflowExecution
 from temporalio.api.enums.v1 import IndexedValueType
@@ -14,11 +14,12 @@
 )
 from temporalio.api.update.v1 import UpdateRef
 from temporalio.api.workflowservice.v1 import PollWorkflowExecutionUpdateRequest
-from temporalio.client import BuildIdOpAddNewDefault, Client
+from temporalio.client import BuildIdOpAddNewDefault, Client, WorkflowHandle
 from temporalio.common import SearchAttributeKey
 from temporalio.service import RPCError, RPCStatusCode
 from temporalio.worker import Worker, WorkflowRunner
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
+from temporalio.workflow import UpdateMethodMultiParam
 
 
 def new_worker(
@@ -128,3 +129,24 @@ async def workflow_update_exists(
         if err.status != RPCStatusCode.NOT_FOUND:
             raise
         return False
+
+
+# TODO: type update return value
+async def admitted_update_task(
+    client: Client,
+    handle: WorkflowHandle,
+    update_method: UpdateMethodMultiParam,
+    id: str,
+    **kwargs,
+) -> asyncio.Task:
+    """
+    Return an asyncio.Task for an update after waiting for it to be admitted.
+    """
+    update_task = asyncio.create_task(
+        handle.execute_update(update_method, id=id, **kwargs)
+    )
+    await assert_eq_eventually(
+        True,
+        lambda: workflow_update_exists(client, handle.id, id),
+    )
+    return update_task