from __future__ import annotations

import asyncio

import concurrent.futures

import dataclasses

import json

import logging

import logging.handlers

import os

import queue

import random

import sys

import threading

import time

import typing

import uuid

from abc import ABC, abstractmethod

from contextlib import contextmanager

from dataclasses import dataclass

from datetime import datetime, timedelta, timezone

from enum import IntEnum

from functools import partial

from typing import (

Any,

Awaitable,

Callable,

Dict,

List,

Mapping,

NoReturn,

Optional,

Sequence,

Tuple,

Type,

Union,

cast,

)

from urllib.request import urlopen

import pydantic

from google.protobuf.timestamp_pb2 import Timestamp

from typing_extensions import Literal, Protocol, runtime_checkable

import temporalio.activity

import temporalio.worker

import temporalio.workflow

from temporalio import activity, workflow

from temporalio.api.common.v1 import Payload, Payloads, WorkflowExecution

from temporalio.api.enums.v1 import EventType

from temporalio.api.failure.v1 import Failure

from temporalio.api.sdk.v1 import EnhancedStackTrace

from temporalio.api.workflowservice.v1 import (

GetWorkflowExecutionHistoryRequest,

ResetStickyTaskQueueRequest,

)

from temporalio.bridge.proto.workflow_activation import WorkflowActivation

from temporalio.bridge.proto.workflow_completion import WorkflowActivationCompletion

from temporalio.client import (

AsyncActivityCancelledError,

Client,

RPCError,

RPCStatusCode,

WorkflowExecutionStatus,

WorkflowFailureError,

WorkflowHandle,

WorkflowQueryFailedError,

WorkflowUpdateFailedError,

WorkflowUpdateHandle,

WorkflowUpdateRPCTimeoutOrCancelledError,

WorkflowUpdateStage,

)

from temporalio.common import (

Priority,

RawValue,

RetryPolicy,

SearchAttributeKey,

SearchAttributePair,

SearchAttributes,

SearchAttributeValues,

TypedSearchAttributes,

WorkflowIDConflictPolicy,

)

from temporalio.converter import (

DataConverter,

DefaultFailureConverter,

DefaultFailureConverterWithEncodedAttributes,

DefaultPayloadConverter,

PayloadCodec,

PayloadConverter,

)

from temporalio.exceptions import (

ActivityError,

ApplicationError,

ApplicationErrorCategory,

CancelledError,

ChildWorkflowError,

TemporalError,

TimeoutError,

WorkflowAlreadyStartedError,

)

from temporalio.runtime import (

BUFFERED_METRIC_KIND_COUNTER,

BUFFERED_METRIC_KIND_HISTOGRAM,

MetricBuffer,

MetricBufferDurationFormat,

PrometheusConfig,

Runtime,

TelemetryConfig,

)

from temporalio.service import __version__

from temporalio.testing import WorkflowEnvironment

from temporalio.worker import (

UnsandboxedWorkflowRunner,

Worker,

WorkflowInstance,

WorkflowInstanceDetails,

WorkflowRunner,

)

from tests import DEV_SERVER_DOWNLOAD_VERSION

from tests.helpers import (

admitted_update_task,

assert_eq_eventually,

assert_eventually,

assert_pending_activity_exists_eventually,

assert_task_fail_eventually,

assert_workflow_exists_eventually,

ensure_search_attributes_present,

find_free_port,

get_pending_activity_info,

new_worker,

pause_and_assert,

unpause_and_assert,

workflow_update_exists,

)

from tests.helpers.external_stack_trace import (

ExternalStackTraceWorkflow,

external_wait_cancel,

)

# Passing through because Python 3.9 has an import bug at

# https://github.com/python/cpython/issues/91351

with workflow.unsafe.imports_passed_through():

import pytest

@workflow.defn

class HelloWorkflow:

@workflow.run

async def run(self, name: str) -> str:

return f"Hello, {name}!"

async def test_workflow_hello(client: Client):

async with new_worker(client, HelloWorkflow) as worker:

result = await client.execute_workflow(

HelloWorkflow.run,

"Temporal",

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

assert result == "Hello, Temporal!"

async def test_workflow_hello_eager(client: Client):

async with new_worker(client, HelloWorkflow) as worker:

handle = await client.start_workflow(

HelloWorkflow.run,

"Temporal",

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

request_eager_start=True,

task_timeout=timedelta(hours=1), # hang if retry needed

)

assert handle.__temporal_eagerly_started

result = await handle.result()

assert result == "Hello, Temporal!"

@activity.defn

async def multi_param_activity(param1: int, param2: str) -> str:

return f"param1: {param1}, param2: {param2}"

@workflow.defn

class MultiParamWorkflow:

@workflow.run

async def run(self, param1: int, param2: str) -> str:

return await workflow.execute_activity(

multi_param_activity,

args=[param1, param2],

schedule_to_close_timeout=timedelta(seconds=30),

)

async def test_workflow_multi_param(client: Client):

# This test is mostly just here to confirm MyPy type checks the multi-param

# overload approach properly

async with new_worker(

client, MultiParamWorkflow, activities=[multi_param_activity]

) as worker:

result = await client.execute_workflow(

MultiParamWorkflow.run,

args=[123, "val1"],

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

assert result == "param1: 123, param2: val1"

@workflow.defn

class InfoWorkflow:

@workflow.run

async def run(self) -> Dict:

# Convert to JSON and back so it'll stringify un-JSON-able pieces

ret = dataclasses.asdict(workflow.info())

return json.loads(json.dumps(ret, default=str))

async def test_workflow_info(client: Client, env: WorkflowEnvironment):

# TODO(cretz): Fix

if env.supports_time_skipping:

pytest.skip(

"Java test server: https://github.com/temporalio/sdk-java/issues/1426"

)

async with new_worker(client, InfoWorkflow) as worker:

workflow_id = f"workflow-{uuid.uuid4()}"

retry_policy = RetryPolicy(

initial_interval=timedelta(seconds=3),

backoff_coefficient=4.0,

maximum_interval=timedelta(seconds=5),

maximum_attempts=6,

)

handle = await client.start_workflow(

InfoWorkflow.run,

id=workflow_id,

task_queue=worker.task_queue,

retry_policy=retry_policy,

)

info = await handle.result()

assert info["attempt"] == 1

assert info["cron_schedule"] is None

assert info["execution_timeout"] is None

assert info["namespace"] == client.namespace

assert info["retry_policy"] == json.loads(

json.dumps(dataclasses.asdict(retry_policy), default=str)

)

assert uuid.UUID(info["run_id"]).version == 7

assert info["run_timeout"] is None

assert info["task_queue"] == worker.task_queue

assert info["task_timeout"] == "0:00:10"

assert info["workflow_id"] == workflow_id

assert info["workflow_type"] == "InfoWorkflow"

async for e in handle.fetch_history_events():

if e.HasField("workflow_execution_started_event_attributes"):

assert info["workflow_start_time"] == json.loads(

json.dumps(

e.event_time.ToDatetime().replace(tzinfo=timezone.utc),

default=str,

)

)

elif e.HasField("workflow_task_started_event_attributes"):

assert info["start_time"] == json.loads(

json.dumps(

e.event_time.ToDatetime().replace(tzinfo=timezone.utc),

default=str,

)

)

@dataclass

class HistoryInfo:

history_length: int

history_size: int

continue_as_new_suggested: bool

@workflow.defn

class HistoryInfoWorkflow:

@workflow.run

async def run(self) -> None:

# Just wait forever

await workflow.wait_condition(lambda: False)

@workflow.signal

async def bunch_of_events(self, count: int) -> None:

# Create a lot of one-day timers

for _ in range(count):

asyncio.create_task(asyncio.sleep(60 * 60 * 24))

@workflow.query

def get_history_info(self) -> HistoryInfo:

return HistoryInfo(

history_length=workflow.info().get_current_history_length(),

history_size=workflow.info().get_current_history_size(),

continue_as_new_suggested=workflow.info().is_continue_as_new_suggested(),

)

async def test_workflow_history_info(

client: Client, env: WorkflowEnvironment, continue_as_new_suggest_history_count: int

):

if env.supports_time_skipping:

pytest.skip("Java test server does not support should continue as new")

async with new_worker(client, HistoryInfoWorkflow) as worker:

handle = await client.start_workflow(

HistoryInfoWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

# Issue query before anything else, which should mean only a history

# size of 3, at least 100 bytes of history, and no continue as new

# suggestion

orig_info = await handle.query(HistoryInfoWorkflow.get_history_info)

assert orig_info.history_length == 3

assert orig_info.history_size > 100

assert not orig_info.continue_as_new_suggested

# Now send a lot of events

await handle.signal(

HistoryInfoWorkflow.bunch_of_events, continue_as_new_suggest_history_count

)

# Send one more event to trigger the WFT update. We have to do this

# because just a query will have a stale representation of history

# counts, but signal forces a new WFT.

await handle.signal(HistoryInfoWorkflow.bunch_of_events, 1)

new_info = await handle.query(HistoryInfoWorkflow.get_history_info)

assert new_info.history_length > continue_as_new_suggest_history_count

assert new_info.history_size > orig_info.history_size

assert new_info.continue_as_new_suggested

@workflow.defn

class SignalAndQueryWorkflow:

def __init__(self) -> None:

self._last_event: Optional[str] = None

@workflow.run

async def run(self) -> None:

# Wait forever

await asyncio.Future()

@workflow.signal

def signal1(self, arg: str) -> None:

self._last_event = f"signal1: {arg}"

@workflow.signal(dynamic=True)

def signal_dynamic(self, name: str, args: Sequence[RawValue]) -> None:

arg = workflow.payload_converter().from_payload(args[0].payload, str)

self._last_event = f"signal_dynamic {name}: {arg}"

@workflow.signal(name="Custom Name")

def signal_custom(self, arg: str) -> None:

self._last_event = f"signal_custom: {arg}"

@workflow.query

def last_event(self) -> str:

return self._last_event or "<no event>"

@workflow.query(dynamic=True)

def query_dynamic(self, name: str, args: Sequence[RawValue]) -> str:

arg = workflow.payload_converter().from_payload(args[0].payload, str)

return f"query_dynamic {name}: {arg}"

@workflow.query(name="Custom Name")

def query_custom(self, arg: str) -> str:

return f"query_custom: {arg}"

async def test_workflow_signal_and_query(client: Client):

async with new_worker(client, SignalAndQueryWorkflow) as worker:

handle = await client.start_workflow(

SignalAndQueryWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

# Simple signals and queries

await handle.signal(SignalAndQueryWorkflow.signal1, "some arg")

assert "signal1: some arg" == await handle.query(

SignalAndQueryWorkflow.last_event

)

# Dynamic signals and queries (old form)

await handle.signal("signal2", "dyn arg")

assert "signal_dynamic signal2: dyn arg" == await handle.query(

SignalAndQueryWorkflow.last_event

)

assert "query_dynamic query2: dyn arg" == await handle.query(

"query2", "dyn arg"

)

# Custom named signals and queries

await handle.signal("Custom Name", "custom arg1")

assert "signal_custom: custom arg1" == await handle.query(

SignalAndQueryWorkflow.last_event

)

await handle.signal(SignalAndQueryWorkflow.signal_custom, "custom arg2")

assert "signal_custom: custom arg2" == await handle.query(

SignalAndQueryWorkflow.last_event

)

assert "query_custom: custom arg1" == await handle.query(

"Custom Name", "custom arg1"

)

assert "query_custom: custom arg1" == await handle.query(

SignalAndQueryWorkflow.query_custom, "custom arg1"

)

@workflow.defn

class SignalAndQueryHandlersWorkflow:

def __init__(self) -> None:

self._last_event: Optional[str] = None

@workflow.run

async def run(self) -> None:

# Wait forever

await asyncio.Future()

@workflow.query

def last_event(self) -> str:

return self._last_event or "<no event>"

@workflow.signal

def set_signal_handler(self, signal_name: str) -> None:

def new_handler(arg: str) -> None:

self._last_event = f"signal {signal_name}: {arg}"

workflow.set_signal_handler(signal_name, new_handler)

@workflow.signal

def set_query_handler(self, query_name: str) -> None:

def new_handler(arg: str) -> str:

return f"query {query_name}: {arg}"

workflow.set_query_handler(query_name, new_handler)

@workflow.signal

def set_dynamic_signal_handler(self) -> None:

def new_handler(name: str, args: Sequence[RawValue]) -> None:

arg = workflow.payload_converter().from_payload(args[0].payload, str)

self._last_event = f"signal dynamic {name}: {arg}"

workflow.set_dynamic_signal_handler(new_handler)

@workflow.signal

def set_dynamic_query_handler(self) -> None:

def new_handler(name: str, args: Sequence[RawValue]) -> str:

arg = workflow.payload_converter().from_payload(args[0].payload, str)

return f"query dynamic {name}: {arg}"

workflow.set_dynamic_query_handler(new_handler)

async def test_workflow_signal_and_query_handlers(client: Client):

async with new_worker(client, SignalAndQueryHandlersWorkflow) as worker:

handle = await client.start_workflow(

SignalAndQueryHandlersWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

# Confirm signals buffered when not found

await handle.signal("unknown_signal1", "val1")

await handle.signal(

SignalAndQueryHandlersWorkflow.set_signal_handler, "unknown_signal1"

)

assert "signal unknown_signal1: val1" == await handle.query(

SignalAndQueryHandlersWorkflow.last_event

)

# Normal signal handling

await handle.signal("unknown_signal1", "val2")

assert "signal unknown_signal1: val2" == await handle.query(

SignalAndQueryHandlersWorkflow.last_event

)

# Dynamic signal handling buffered and new

await handle.signal("unknown_signal2", "val3")

await handle.signal(SignalAndQueryHandlersWorkflow.set_dynamic_signal_handler)

assert "signal dynamic unknown_signal2: val3" == await handle.query(

SignalAndQueryHandlersWorkflow.last_event

)

await handle.signal("unknown_signal3", "val4")

assert "signal dynamic unknown_signal3: val4" == await handle.query(

SignalAndQueryHandlersWorkflow.last_event

)

# Normal query handling

await handle.signal(

SignalAndQueryHandlersWorkflow.set_query_handler, "unknown_query1"

)

assert "query unknown_query1: val5" == await handle.query(

"unknown_query1", "val5"

)

# Dynamic query handling

await handle.signal(SignalAndQueryHandlersWorkflow.set_dynamic_query_handler)

assert "query dynamic unknown_query2: val6" == await handle.query(

"unknown_query2", "val6"

)

@workflow.defn

class SignalAndQueryErrorsWorkflow:

@workflow.run

async def run(self) -> None:

# Wait forever

await asyncio.Future()

@workflow.signal

def bad_signal(self) -> NoReturn:

raise ApplicationError("signal fail", 123)

@workflow.query

def bad_query(self) -> NoReturn:

raise ApplicationError("query fail", 456)

@workflow.query

def other_query(self) -> str:

raise NotImplementedError

async def test_workflow_signal_and_query_errors(client: Client):

async with new_worker(client, SignalAndQueryErrorsWorkflow) as worker:

handle = await client.start_workflow(

SignalAndQueryErrorsWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

# Send bad signal

await handle.signal(SignalAndQueryErrorsWorkflow.bad_signal)

# Wait on workflow

with pytest.raises(WorkflowFailureError) as err:

await handle.result()

assert isinstance(err.value.cause, ApplicationError)

assert list(err.value.cause.details) == [123]

# Fail query (no details on query failure)

with pytest.raises(WorkflowQueryFailedError) as rpc_err:

await handle.query(SignalAndQueryErrorsWorkflow.bad_query)

assert str(rpc_err.value) == "query fail"

# Unrecognized query

with pytest.raises(WorkflowQueryFailedError) as rpc_err:

await handle.query("non-existent query")

assert str(rpc_err.value) == (

"Query handler for 'non-existent query' expected but not found,"

" known queries: [__enhanced_stack_trace __stack_trace __temporal_workflow_metadata bad_query other_query]"

)

@workflow.defn

class SignalAndQueryOldDynamicStyleWorkflow:

def __init__(self) -> None:

self._last_event: Optional[str] = None

@workflow.run

async def run(self) -> None:

# Wait forever

await asyncio.Future()

@workflow.signal(dynamic=True)

def signal_dynamic(self, name: str, *args: Any) -> None:

self._last_event = f"signal_dynamic {name}: {args[0]}"

@workflow.query

def last_event(self) -> str:

return self._last_event or "<no event>"

@workflow.query(dynamic=True)

def query_dynamic(self, name: str, *args: Any) -> str:

return f"query_dynamic {name}: {args[0]}"

async def test_workflow_signal_and_query_old_dynamic_style(client: Client):

async with new_worker(client, SignalAndQueryOldDynamicStyleWorkflow) as worker:

handle = await client.start_workflow(

SignalAndQueryOldDynamicStyleWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

# Dynamic signals and queries

await handle.signal("signal1", "dyn arg")

assert "signal_dynamic signal1: dyn arg" == await handle.query(

SignalAndQueryOldDynamicStyleWorkflow.last_event

)

assert "query_dynamic query1: dyn arg" == await handle.query(

"query1", "dyn arg"

)

@workflow.defn

class SignalAndQueryHandlersOldDynamicStyleWorkflow:

def __init__(self) -> None:

self._last_event: Optional[str] = None

@workflow.run

async def run(self) -> None:

# Wait forever

await asyncio.Future()

@workflow.query

def last_event(self) -> str:

return self._last_event or "<no event>"

@workflow.signal

def set_dynamic_signal_handler(self) -> None:

def new_handler(name: str, *args: Any) -> None:

self._last_event = f"signal dynamic {name}: {args[0]}"

workflow.set_dynamic_signal_handler(new_handler)

@workflow.signal

def set_dynamic_query_handler(self) -> None:

def new_handler(name: str, *args: Any) -> str:

return f"query dynamic {name}: {args[0]}"

workflow.set_dynamic_query_handler(new_handler)

async def test_workflow_signal_qnd_query_handlers_old_dynamic_style(client: Client):

async with new_worker(

client, SignalAndQueryHandlersOldDynamicStyleWorkflow

) as worker:

handle = await client.start_workflow(

SignalAndQueryHandlersOldDynamicStyleWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

# Dynamic signal handling buffered and new

await handle.signal("unknown_signal1", "val1")

await handle.signal(

SignalAndQueryHandlersOldDynamicStyleWorkflow.set_dynamic_signal_handler

)

assert "signal dynamic unknown_signal1: val1" == await handle.query(

SignalAndQueryHandlersOldDynamicStyleWorkflow.last_event

)

await handle.signal("unknown_signal2", "val2")

assert "signal dynamic unknown_signal2: val2" == await handle.query(

SignalAndQueryHandlersOldDynamicStyleWorkflow.last_event

)

# Dynamic query handling

await handle.signal(

SignalAndQueryHandlersOldDynamicStyleWorkflow.set_dynamic_query_handler

)

assert "query dynamic unknown_query1: val3" == await handle.query(

"unknown_query1", "val3"

)

@dataclass

class BadSignalParam:

some_str: str

@workflow.defn

class BadSignalParamWorkflow:

def __init__(self) -> None:

self._signals: List[BadSignalParam] = []

@workflow.run

async def run(self) -> List[BadSignalParam]:

await workflow.wait_condition(

lambda: bool(self._signals) and self._signals[-1].some_str == "finish"

)

return self._signals

@workflow.signal

async def some_signal(self, param: BadSignalParam) -> None:

self._signals.append(param)

async def test_workflow_bad_signal_param(client: Client):

async with new_worker(client, BadSignalParamWorkflow) as worker:

handle = await client.start_workflow(

BadSignalParamWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

# Send 4 signals, first and third are bad

await handle.signal("some_signal", "bad")

await handle.signal("some_signal", BadSignalParam(some_str="good"))

await handle.signal("some_signal", 123)

await handle.signal("some_signal", BadSignalParam(some_str="finish"))

assert [

BadSignalParam(some_str="good"),

BadSignalParam(some_str="finish"),

] == await handle.result()

@workflow.defn

class AsyncUtilWorkflow:

def __init__(self) -> None:

self._status = "starting"

self._wait_event1 = asyncio.Event()

self._received_event2 = False

@workflow.run

async def run(self) -> Dict:

# Record start times

ret = {

# "now" timestamp and current event loop monotonic time

"start": str(workflow.now()),

"start_time": workflow.time(),

"start_time_ns": workflow.time_ns(),

"event_loop_start": asyncio.get_running_loop().time(),

}

# Sleep for a small amount of time (we accept that it may take longer on

# the server)

await asyncio.sleep(0.1)

# Wait for event 1

self._status = "waiting for event1"

await self._wait_event1.wait()

# Wait for event 2

self._status = "waiting for event2"

await workflow.wait_condition(lambda: self._received_event2)

# Record completion times

self._status = "done"

ret["end_time_ns"] = workflow.time_ns()

return ret

@workflow.signal

def event1(self) -> None:

self._wait_event1.set()

@workflow.signal

def event2(self) -> None:

self._received_event2 = True

@workflow.query

def status(self) -> str:

return self._status

async def test_workflow_async_utils(client: Client):

async with new_worker(client, AsyncUtilWorkflow) as worker:

# Start workflow and wait until status is waiting for event 1

handle = await client.start_workflow(

AsyncUtilWorkflow.run,

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

async def status() -> str:

return await handle.query(AsyncUtilWorkflow.status)

await assert_eq_eventually("waiting for event1", status)

# Set event 1 and confirm waiting on event 2

await handle.signal(AsyncUtilWorkflow.event1)

await assert_eq_eventually("waiting for event2", status)

# Set event 2 and get the result and confirm query still works

await handle.signal(AsyncUtilWorkflow.event2)

result = await handle.result()

assert "done" == await status()

# Get the actual start time out of history

resp = await client.workflow_service.get_workflow_execution_history(

GetWorkflowExecutionHistoryRequest(

namespace=client.namespace,

execution=WorkflowExecution(workflow_id=handle.id),

)

)

first_timestamp: Optional[Timestamp] = None

last_timestamp: Optional[Timestamp] = None

for event in resp.history.events:

# Get timestamp from first workflow task started

if event.event_type is EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED:

if not first_timestamp:

first_timestamp = event.event_time

last_timestamp = event.event_time

assert first_timestamp and last_timestamp

# Check the times. We have to ignore type here because typeshed has

# wrong type for Protobuf ToDatetime.

first_timestamp_datetime = first_timestamp.ToDatetime(tzinfo=timezone.utc) # type: ignore

# We take off subsecond because Protobuf rounds nanos

# differently than we do (they round toward zero, we use

# utcfromtimestamp which suffers float precision issues).

assert datetime.fromisoformat(result["start"]).replace(

microsecond=0

) == first_timestamp_datetime.replace(microsecond=0)

assert result["start_time"] == first_timestamp.ToNanoseconds() / 1e9

assert result["start_time_ns"] == first_timestamp.ToNanoseconds()

assert result["event_loop_start"] == result["start_time"]

assert result["start_time_ns"] < result["end_time_ns"]

assert result["end_time_ns"] == last_timestamp.ToNanoseconds()

@activity.defn

async def say_hello(name: str) -> str:

return f"Hello, {name}!"

@workflow.defn

class SimpleActivityWorkflow:

@workflow.run

async def run(self, name: str) -> str:

return await workflow.execute_activity(

say_hello, name, schedule_to_close_timeout=timedelta(seconds=5)

)

async def test_workflow_simple_activity(client: Client):

async with new_worker(

client, SimpleActivityWorkflow, activities=[say_hello]

) as worker:

result = await client.execute_workflow(

SimpleActivityWorkflow.run,

"Temporal",

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

assert result == "Hello, Temporal!"

@workflow.defn

class SimpleLocalActivityWorkflow:

@workflow.run

async def run(self, name: str) -> str:

return await workflow.execute_local_activity(

say_hello, name, schedule_to_close_timeout=timedelta(seconds=5)

)

async def test_workflow_simple_local_activity(client: Client):

async with new_worker(

client, SimpleLocalActivityWorkflow, activities=[say_hello]

) as worker:

result = await client.execute_workflow(

SimpleLocalActivityWorkflow.run,

"Temporal",

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

)

assert result == "Hello, Temporal!"

@activity.defn

async def wait_cancel() -> str:

try:

if activity.info().is_local:

await asyncio.sleep(1000)

else:

while True:

await asyncio.sleep(0.3)

activity.heartbeat()

return "Manually stopped"

except asyncio.CancelledError:

return "Got cancelled error, cancelled? " + str(activity.is_cancelled())

class ActivityWaitCancelNotify:

def __init__(self) -> None:

self.wait_cancel_complete = asyncio.Event()

@activity.defn

async def wait_cancel(self) -> str:

self.wait_cancel_complete.clear()

try:

if activity.info().is_local:

await asyncio.sleep(1000)

else:

while True:

await asyncio.sleep(0.3)

activity.heartbeat()

return "Manually stopped"

except asyncio.CancelledError:

return "Got cancelled error, cancelled? " + str(activity.is_cancelled())

finally:

self.wait_cancel_complete.set()

@dataclass

class CancelActivityWorkflowParams:

cancellation_type: str

local: bool

@workflow.defn

class CancelActivityWorkflow:

def __init__(self) -> None:

self._activity_result = "<none>"

@workflow.run

async def run(self, params: CancelActivityWorkflowParams) -> None:

if params.local:

handle = workflow.start_local_activity_method(

ActivityWaitCancelNotify.wait_cancel,

schedule_to_close_timeout=timedelta(seconds=5),

cancellation_type=workflow.ActivityCancellationType[

params.cancellation_type

],

)

else:

handle = workflow.start_activity_method(

ActivityWaitCancelNotify.wait_cancel,

schedule_to_close_timeout=timedelta(seconds=5),

heartbeat_timeout=timedelta(seconds=1),

cancellation_type=workflow.ActivityCancellationType[

params.cancellation_type

],

)

await asyncio.sleep(0.01)

try:

handle.cancel()

self._activity_result = await handle

except ActivityError as err:

self._activity_result = f"Error: {err.cause.__class__.__name__}"

# TODO(cretz): Remove when https://github.com/temporalio/sdk-core/issues/323 is fixed

except CancelledError as err:

self._activity_result = f"Error: {err.__class__.__name__}"

# Wait forever

await asyncio.Future()

@workflow.query

def activity_result(self) -> str:

return self._activity_result

@pytest.mark.parametrize("local", [True, False])

async def test_workflow_cancel_activity(client: Client, local: bool):

# Need short task timeout to timeout LA task and longer assert timeout

# so the task can timeout

task_timeout = timedelta(seconds=1)

assert_timeout = timedelta(seconds=10)

activity_inst = ActivityWaitCancelNotify()

async with new_worker(

client, CancelActivityWorkflow, activities=[activity_inst.wait_cancel]

) as worker:

# Try cancel - confirm error and activity was sent the cancel

handle = await client.start_workflow(

CancelActivityWorkflow.run,

CancelActivityWorkflowParams(

cancellation_type=workflow.ActivityCancellationType.TRY_CANCEL.name,

local=local,

),

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

task_timeout=task_timeout,

)

async def activity_result() -> str:

return await handle.query(CancelActivityWorkflow.activity_result)

await assert_eq_eventually(

"Error: CancelledError", activity_result, timeout=assert_timeout

)

await activity_inst.wait_cancel_complete.wait()

await handle.cancel()

# Wait cancel - confirm no error due to graceful cancel handling

handle = await client.start_workflow(

CancelActivityWorkflow.run,

CancelActivityWorkflowParams(

cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED.name,

local=local,

),

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

task_timeout=task_timeout,

)

await assert_eq_eventually(

"Got cancelled error, cancelled? True",

activity_result,

timeout=assert_timeout,

)

await activity_inst.wait_cancel_complete.wait()

await handle.cancel()

# Abandon - confirm error and that activity stays running

handle = await client.start_workflow(

CancelActivityWorkflow.run,

CancelActivityWorkflowParams(

cancellation_type=workflow.ActivityCancellationType.ABANDON.name,

local=local,

),

id=f"workflow-{uuid.uuid4()}",

task_queue=worker.task_queue,

task_timeout=task_timeout,

)

await assert_eq_eventually(

"Error: CancelledError", activity_result, timeout=assert_timeout

)

await asyncio.sleep(0.5)

assert not activity_inst.wait_cancel_complete.is_set()

await handle.cancel()

await activity_inst.wait_cancel_complete.wait()

@workflow.defn

class SimpleChildWorkflow: