84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
from typing import Any, Optional
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
|
|
from workflows import Workflow, Context
|
|
from workflows.events import StartEvent, Event, EventType
|
|
|
|
|
|
@dataclass
|
|
class WorkflowTestResult:
|
|
"""
|
|
Container for workflow test results
|
|
|
|
Attributes:
|
|
collected (list[Event]): List of collected events
|
|
event_type (dict[EventType, int]): Dictionary that maps each event type with its number of occurencies within the collected events
|
|
result (Any): Final output of the workflow run
|
|
"""
|
|
|
|
collected: list[Event]
|
|
event_types: dict[EventType, int]
|
|
result: Any
|
|
ctx: Context
|
|
|
|
|
|
class WorkflowTestRunner:
|
|
"""
|
|
Utility class that can be used to test workflows end-to-end.
|
|
|
|
Attributes:
|
|
_workflow (Workflow): The workflow to be tested
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
workflow: "Workflow",
|
|
):
|
|
self._workflow = workflow
|
|
|
|
async def run(
|
|
self,
|
|
start_event: StartEvent = StartEvent(),
|
|
ctx: Optional["Context"] = None,
|
|
expose_internal: bool = True,
|
|
exclude_events: Optional[list[EventType]] = None,
|
|
) -> WorkflowTestResult:
|
|
"""
|
|
Run a workflow end-to-end and collect the events that are streamed during its execution.
|
|
|
|
Args:
|
|
start_event (StartEvent): The input event for the workflow
|
|
expose_internal (bool): Whether or not to expose internal events. Defaults to True if not set.
|
|
exclude_events. (list[EventType]): A list of event types to exclude from the collected events. Defaults to None if not set.
|
|
|
|
Returns:
|
|
WorkflowTestResult
|
|
|
|
Example:
|
|
```
|
|
wf = GreetingWorkflow()
|
|
runner = WorkflowTestRunner(wf)
|
|
test_result = runner.run(start_even=StartEvent(message="hello"), expose_internal = True, exclude_events = [StepStateChanged])
|
|
assert test_result.collected == 22
|
|
assert test_result.event_types.get(StepStateChanged, 0) == 8
|
|
assert str(test_result.result) == "hello Adam!"
|
|
```
|
|
"""
|
|
handler = self._workflow.run(start_event=start_event, ctx=ctx)
|
|
collected_events: list[Event] = []
|
|
async for event in handler.stream_events(expose_internal=expose_internal):
|
|
if exclude_events and type(event) in exclude_events:
|
|
continue
|
|
collected_events.append(event)
|
|
result = await handler
|
|
event_freqs: dict[EventType, int] = dict(
|
|
Counter([type(ev) for ev in collected_events])
|
|
)
|
|
assert handler.ctx is not None
|
|
return WorkflowTestResult(
|
|
collected=collected_events,
|
|
result=result,
|
|
event_types=event_freqs,
|
|
ctx=handler.ctx,
|
|
)
|