@workflow.defn
class UpdateCancellationWorkflow:
    def __init__(self) -> None:
        self.non_terminating_operation_has_started = False

    @workflow.run
    async def run(self) -> str:
        await workflow.wait_condition(lambda: False)
        return "unreachable"

    @workflow.update
    async def wait_until_non_terminating_operation_has_started(self) -> None:
        await workflow.wait_condition(
            lambda: self.non_terminating_operation_has_started
        )

    @workflow.update
    async def non_terminating_operation(self) -> str:
        self.non_terminating_operation_has_started = True
        await workflow.wait_condition(lambda: False)
        return "unreachable"


async def test_update_cancellation(client: Client):
    async with new_worker(client, UpdateCancellationWorkflow) as worker:
        wf_handle = await client.start_workflow(
            UpdateCancellationWorkflow.run,
            id=str(uuid.uuid4()),
            task_queue=worker.task_queue,
        )
        # Asynchronously run an update that will never complete
        non_terminating_update_task = asyncio.create_task(
            wf_handle.execute_update(
                UpdateCancellationWorkflow.non_terminating_operation
            )
        )
        print("wait until handler started...")
        # Wait until we know the update handler has started executing
        await wf_handle.execute_update(
            UpdateCancellationWorkflow.wait_until_non_terminating_operation_has_started
        )
        print("cancel the workflow")
        await wf_handle.cancel()
        print("await non_terminating_update_task...")
        await non_terminating_update_task