import asyncio

import uuid

from datetime import timedelta

from typing import Any, Callable, List, NoReturn, Optional, Tuple, Type

import pytest

from temporalio import activity, workflow

from temporalio.client import Client, WorkflowUpdateFailedError

from temporalio.exceptions import ApplicationError

from temporalio.testing import WorkflowEnvironment

from temporalio.worker import (

ActivityInboundInterceptor,

ActivityOutboundInterceptor,

ContinueAsNewInput,

ExecuteActivityInput,

ExecuteWorkflowInput,

HandleQueryInput,

HandleSignalInput,

HandleUpdateInput,

Interceptor,

SignalChildWorkflowInput,

SignalExternalWorkflowInput,

StartActivityInput,

StartChildWorkflowInput,

StartLocalActivityInput,

Worker,

WorkflowInboundInterceptor,

WorkflowInterceptorClassInput,

WorkflowOutboundInterceptor,

)

interceptor_traces: List[Tuple[str, Any]] = []

class TracingWorkerInterceptor(Interceptor):

def intercept_activity(

self, next: ActivityInboundInterceptor

) -> ActivityInboundInterceptor:

return TracingActivityInboundInterceptor(super().intercept_activity(next))

def workflow_interceptor_class(

self, input: WorkflowInterceptorClassInput

) -> Optional[Type[WorkflowInboundInterceptor]]:

return TracingWorkflowInboundInterceptor

class TracingActivityInboundInterceptor(ActivityInboundInterceptor):

def init(self, outbound: ActivityOutboundInterceptor) -> None:

super().init(TracingActivityOutboundInterceptor(outbound))

async def execute_activity(self, input: ExecuteActivityInput) -> Any:

interceptor_traces.append(("activity.execute", input))

return await super().execute_activity(input)

class TracingActivityOutboundInterceptor(ActivityOutboundInterceptor):

def info(self) -> activity.Info:

interceptor_traces.append(("activity.info", super().info()))

return super().info()

def heartbeat(self, *details: Any) -> None:

interceptor_traces.append(("activity.heartbeat", details))

super().heartbeat(*details)

class TracingWorkflowInboundInterceptor(WorkflowInboundInterceptor):

def init(self, outbound: WorkflowOutboundInterceptor) -> None:

super().init(TracingWorkflowOutboundInterceptor(outbound))

async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:

interceptor_traces.append(("workflow.execute", input))

return await super().execute_workflow(input)

async def handle_signal(self, input: HandleSignalInput) -> None:

interceptor_traces.append(("workflow.signal", input))

return await super().handle_signal(input)

async def handle_query(self, input: HandleQueryInput) -> Any:

interceptor_traces.append(("workflow.query", input))

return await super().handle_query(input)

def handle_update_validator(self, input: HandleUpdateInput) -> None:

interceptor_traces.append(("workflow.update.validator", input))

return super().handle_update_validator(input)

async def handle_update_handler(self, input: HandleUpdateInput) -> Any:

interceptor_traces.append(("workflow.update.handler", input))

return await super().handle_update_handler(input)

class TracingWorkflowOutboundInterceptor(WorkflowOutboundInterceptor):

def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn:

interceptor_traces.append(("workflow.continue_as_new", input))

super().continue_as_new(input)

def info(self) -> workflow.Info:

interceptor_traces.append(("workflow.info", super().info()))

return super().info()

async def signal_child_workflow(self, input: SignalChildWorkflowInput) -> None:

interceptor_traces.append(("workflow.signal_child_workflow", input))

await super().signal_child_workflow(input)

async def signal_external_workflow(

self, input: SignalExternalWorkflowInput

) -> None:

interceptor_traces.append(("workflow.signal_external_workflow", input))

await super().signal_external_workflow(input)

def start_activity(self, input: StartActivityInput) -> workflow.ActivityHandle:

interceptor_traces.append(("workflow.start_activity", input))

return super().start_activity(input)

async def start_child_workflow(

self, input: StartChildWorkflowInput

) -> workflow.ChildWorkflowHandle:

interceptor_traces.append(("workflow.start_child_workflow", input))

return await super().start_child_workflow(input)

def start_local_activity(

self, input: StartLocalActivityInput

) -> workflow.ActivityHandle:

interceptor_traces.append(("workflow.start_local_activity", input))

return super().start_local_activity(input)

@activity.defn

async def intercepted_activity(param: str) -> str:

if not activity.info().is_local:

activity.heartbeat("details")

return f"param: {param}"

@workflow.defn

class InterceptedWorkflow:

def __init__(self) -> None:

self.finish = asyncio.Event()

@workflow.run

async def run(self, style: str) -> None:

if style == "continue-as-new":

return

if style == "child" or style == "external":

await self.finish.wait()

return

await workflow.execute_activity(

intercepted_activity, "val1", schedule_to_close_timeout=timedelta(seconds=5)

)

await workflow.execute_local_activity(

intercepted_activity, "val2", schedule_to_close_timeout=timedelta(seconds=5)

)

my_id = workflow.info().workflow_id

child_handle = await workflow.start_child_workflow(

InterceptedWorkflow.run, "child", id=f"{my_id}_child"

)

await child_handle.signal(InterceptedWorkflow.signal, "child-signal-val")

await child_handle

# Create another child so we can use it for external handle

child_handle = await workflow.start_child_workflow(

InterceptedWorkflow.run, "external", id=f"{my_id}_external"

)

await workflow.get_external_workflow_handle(child_handle.id).signal(

InterceptedWorkflow.signal, "external-signal-val"

)

await child_handle

await self.finish.wait()

workflow.continue_as_new("continue-as-new")

@workflow.query

def query(self, param: str) -> str:

return f"query: {param}"

@workflow.signal

def signal(self, param: str) -> None:

self.finish.set()

@workflow.update

def update(self, param: str) -> str:

return f"update: {param}"

@workflow.update

def update_validated(self, param: str) -> str:

return f"update: {param}"

@update_validated.validator

def update_validated_validator(self, param: str) -> None:

if param == "reject-me":

raise ApplicationError("Invalid update")

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

# TODO(cretz): Fix

if env.supports_time_skipping:

pytest.skip(

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

)

task_queue = f"task_queue_{uuid.uuid4()}"

async with Worker(

client,

task_queue=task_queue,

workflows=[InterceptedWorkflow],

activities=[intercepted_activity],

interceptors=[TracingWorkerInterceptor()],

):

# Run workflow

handle = await client.start_workflow(

InterceptedWorkflow.run,

"initial",

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

task_queue=task_queue,

)

assert "query: query-val" == await handle.query(

InterceptedWorkflow.query, "query-val"

)

await handle.signal(InterceptedWorkflow.signal, "signal-val")

assert "update: update-val" == await handle.execute_update(

InterceptedWorkflow.update, "update-val"

)

with pytest.raises(WorkflowUpdateFailedError) as _err:

await handle.execute_update(

InterceptedWorkflow.update_validated, "reject-me"

)

await handle.result()

# Check traces

def pop_trace(name: str, filter: Optional[Callable[[Any], bool]] = None) -> Any:

index = next(

(

i

for i, v in enumerate(interceptor_traces)

if v[0] == name and (not filter or filter(v[1]))

),

None,

)

if index is None:

return None

return interceptor_traces.pop(index)[1]

assert pop_trace("activity.execute", lambda v: v.args[0] == "val1")

assert pop_trace("activity.execute", lambda v: v.args[0] == "val2")

# Check activity info is called _at least_ twice, but is called more

# because the logger uses it for context

activity_infos = 0

while pop_trace("activity.info"):

activity_infos += 1

assert activity_infos >= 2

assert pop_trace("activity.heartbeat", lambda v: v[0] == "details")

# One initial, one child, one external, one continue as new

assert pop_trace("workflow.execute", lambda v: v.args[0] == "initial")

assert pop_trace("workflow.execute", lambda v: v.args[0] == "child")

assert pop_trace("workflow.execute", lambda v: v.args[0] == "external")

assert pop_trace("workflow.execute", lambda v: v.args[0] == "continue-as-new")

assert pop_trace("workflow.signal", lambda v: v.args[0] == "signal-val")

assert pop_trace("workflow.query", lambda v: v.args[0] == "query-val")

assert pop_trace("workflow.continue_as_new")

assert pop_trace("workflow.info")

assert pop_trace("workflow.start_activity", lambda v: v.args[0] == "val1")

assert pop_trace("workflow.start_local_activity", lambda v: v.args[0] == "val2")

assert pop_trace(

"workflow.start_child_workflow", lambda v: v.args[0] == "child"

)

assert pop_trace(

"workflow.signal_child_workflow", lambda v: v.args[0] == "child-signal-val"

)

assert pop_trace("workflow.signal", lambda v: v.args[0] == "child-signal-val")

assert pop_trace(

"workflow.start_child_workflow", lambda v: v.args[0] == "external"

)

assert pop_trace(

"workflow.signal_external_workflow",

lambda v: v.args[0] == "external-signal-val",

)

assert pop_trace(

"workflow.signal", lambda v: v.args[0] == "external-signal-val"

)

assert pop_trace("workflow.update.handler", lambda v: v.args[0] == "update-val")

assert pop_trace(

"workflow.update.validator", lambda v: v.args[0] == "reject-me"

)

# Confirm no unexpected traces

assert not interceptor_traces