From d36066fcdd9e6600e35016641275a12d403f7288 Mon Sep 17 00:00:00 2001 From: jcschaff Date: Mon, 13 Jan 2025 23:33:06 -0500 Subject: [PATCH] add pydantic Temporal test --- tests/workflows/test_temporal.py | 35 ++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/workflows/test_temporal.py b/tests/workflows/test_temporal.py index b0d5772..89643f3 100644 --- a/tests/workflows/test_temporal.py +++ b/tests/workflows/test_temporal.py @@ -1,11 +1,11 @@ -import pytest import asyncio from datetime import timedelta from typing import List +import pytest +from pydantic import BaseModel from temporalio import activity, workflow from temporalio.client import Client - from temporalio.worker import Worker, UnsandboxedWorkflowRunner @@ -59,3 +59,34 @@ async def test_simple_workflow(temporal_client: Client) -> None: ) expected_results = ['Hello, user1!', 'Hello, user2!', 'Hello, user3!', 'Hello, user4!', 'Hello, user5!'] assert result == expected_results + + + +class User(BaseModel): + name: str + age: int + + +@activity.defn +async def get_hello_pydantic_user_data() -> User: + return User(name="John Doe", age=30) + + +@workflow.defn +class HelloPydanticWorkflow: + @workflow.run + async def run(self) -> User: + return await workflow.execute_activity(get_hello_pydantic_user_data, args=[], + start_to_close_timeout=timedelta(seconds=5)) + + +@pytest.mark.asyncio +async def test_pydantic_workflow(temporal_client: Client) -> None: + # Run a worker for the workflow + async with Worker(temporal_client, task_queue="hello-parallel-activity-task-queue", + workflows=[HelloPydanticWorkflow], activities=[get_hello_pydantic_user_data], + workflow_runner=UnsandboxedWorkflowRunner()): + result = await temporal_client.execute_workflow(HelloPydanticWorkflow.run, args=[], + id="hello-pydantic-workflow-id", task_queue="hello-parallel-activity-task-queue", ) + expected_results = User(name="John Doe", age=30) + assert result == expected_results