81 lines
2.1 KiB
Python
81 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from wf_authoring import build_async_registry, node
|
|
from wf_core import RunStatus, RuntimeContext, execute_workflow_async
|
|
|
|
from .test_demo_workflow import (
|
|
build_authoring_demo_workflow,
|
|
combine_summaries_spec,
|
|
drive_list_files_spec,
|
|
mark_email_skipped_spec,
|
|
send_email_spec,
|
|
summarize_document_spec,
|
|
)
|
|
|
|
|
|
class InferredEchoInput(BaseModel):
|
|
value: str
|
|
|
|
|
|
class InferredEchoOutput(BaseModel):
|
|
echoed: str
|
|
|
|
|
|
class InferredAsyncInput(BaseModel):
|
|
value: str
|
|
|
|
|
|
class InferredAsyncOutput(BaseModel):
|
|
echoed: str
|
|
|
|
|
|
async def test_async_registry_accepts_sync_and_async_specs() -> None:
|
|
@node()
|
|
def sync_echo(
|
|
payload: InferredEchoInput,
|
|
ctx: RuntimeContext,
|
|
) -> InferredEchoOutput:
|
|
return InferredEchoOutput(echoed=payload.value)
|
|
|
|
@node()
|
|
async def async_echo(
|
|
payload: InferredAsyncInput,
|
|
ctx: RuntimeContext,
|
|
) -> InferredAsyncOutput:
|
|
return InferredAsyncOutput(echoed=f"async:{payload.value}")
|
|
|
|
registry = build_async_registry(sync_echo, async_echo)
|
|
ctx = RuntimeContext(current_node_id="demo")
|
|
|
|
async def run_handler(name: str, value: str) -> dict[str, object]:
|
|
return await registry[name]({"value": value}, ctx)
|
|
|
|
sync_result = await run_handler("sync_echo", "hello")
|
|
async_result = await run_handler("async_echo", "world")
|
|
|
|
assert sync_result == {"outcome": "ok", "output": {"echoed": "hello"}}
|
|
assert async_result == {"outcome": "ok", "output": {"echoed": "async:world"}}
|
|
|
|
|
|
async def test_execute_workflow_async_runs_with_async_registry() -> None:
|
|
workflow, _ = build_authoring_demo_workflow()
|
|
registry = build_async_registry(
|
|
drive_list_files_spec,
|
|
summarize_document_spec,
|
|
combine_summaries_spec,
|
|
send_email_spec,
|
|
mark_email_skipped_spec,
|
|
)
|
|
|
|
run = await execute_workflow_async(
|
|
workflow,
|
|
{"folder_id": "demo-folder", "should_email": False},
|
|
registry,
|
|
)
|
|
|
|
assert run.status == RunStatus.COMPLETED
|
|
assert run.output["email_status"] == "skipped"
|
|
assert run.interrupt is None
|