Add tests for all LangGraph samples and fix serialization bugs

Tests:
- Added tests for all LangGraph samples (hello_world, approval_workflow,
  react_agent, supervisor, agentic_rag, deep_research, plan_and_execute,
  reflection)
- Added conftest.py with shared fixtures (clear_registry, requires_openai)
- Tests requiring OpenAI API are skipped when OPENAI_API_KEY is not set

Bug fixes:
- Fixed serialization bugs in reflection and plan_and_execute samples
- After Temporal serialization, Pydantic models become dicts
- Added helper functions to handle both object and dict access for:
  - Critique objects in reflection sample
  - Plan, PlanStep, StepResult objects in plan_and_execute sample

Author: mfateev

langgraph_samples/plan_and_execute/graph.py

@@ -84,6 +84,30 @@ class PlanExecuteState(TypedDict):
     needs_replan: bool
 
 
+# Helper functions to handle dict/object access after Temporal serialization
+def _get_plan_steps(plan: Plan | dict[str, Any]) -> list[Any]:
+    """Get steps from Plan, handling both object and dict forms."""
+    if isinstance(plan, dict):
+        return plan.get("steps", [])
+    return list(plan.steps)
+
+
+def _get_step_attr(
+    step: PlanStep | dict[str, Any], attr: str, default: Any = ""
+) -> Any:
+    """Get attribute from PlanStep, handling both object and dict forms."""
+    if isinstance(step, dict):
+        return step.get(attr, default)
+    return getattr(step, attr, default)
+
+
+def _get_result_success(result: StepResult | dict[str, Any]) -> bool:
+    """Get success from StepResult, handling both object and dict forms."""
+    if isinstance(result, dict):
+        return result.get("success", False)
+    return result.success
+
+
 # Define tools for the executor agent
 @tool
 def calculate(expression: str) -> str:
@@ -220,21 +244,25 @@ def execute_step(state: PlanExecuteState) -> dict[str, Any]:
         plan = state.get("plan")
         current_step = state.get("current_step", 0)
 
-        if not plan or current_step >= len(plan.steps):
+        steps = _get_plan_steps(plan) if plan else []
+        if not plan or current_step >= len(steps):
             return {"needs_replan": False}
 
-        step = plan.steps[current_step]
+        step = steps[current_step]
+        step_number = _get_step_attr(step, "step_number", 0)
+        step_desc = _get_step_attr(step, "description", "")
+        tool_hint = _get_step_attr(step, "tool_hint", "")
 
         # Run the executor agent for this step
         result = executor_agent.invoke(
             {
                 "messages": [
                     SystemMessage(
-                        content=f"You are executing step {step.step_number} of a plan. "
-                        f"Complete this step: {step.description}\n"
-                        f"Suggested tool: {step.tool_hint}"
+                        content=f"You are executing step {step_number} of a plan. "
+                        f"Complete this step: {step_desc}\n"
+                        f"Suggested tool: {tool_hint}"
                     ),
-                    HumanMessage(content=step.description),
+                    HumanMessage(content=step_desc),
                 ]
             }
         )
@@ -248,8 +276,8 @@ def execute_step(state: PlanExecuteState) -> dict[str, Any]:
             )
 
         step_result = StepResult(
-            step_number=step.step_number,
-            description=step.description,
+            step_number=step_number,
+            description=step_desc,
             result=result_content,
             success=True,
         )
@@ -260,7 +288,7 @@ def execute_step(state: PlanExecuteState) -> dict[str, Any]:
             "messages": [
                 {
                     "role": "assistant",
-                    "content": f"Step {step.step_number} completed: {result_content[:200]}...",
+                    "content": f"Step {step_number} completed: {result_content[:200]}...",
                 }
             ],
         }
@@ -278,17 +306,18 @@ def evaluate_progress(state: PlanExecuteState) -> dict[str, Any]:
             return {"needs_replan": True}
 
         # Check if all steps completed
-        all_complete = current_step >= len(plan.steps)
+        steps = _get_plan_steps(plan)
+        all_complete = current_step >= len(steps)
 
         # Check if any step failed
-        failed_steps = [r for r in step_results if not r.success]
+        failed_steps = [r for r in step_results if not _get_result_success(r)]
 
         return {
             "needs_replan": len(failed_steps) > 0,
             "messages": [
                 {
                     "role": "assistant",
-                    "content": f"Progress: {current_step}/{len(plan.steps)} steps complete. "
+                    "content": f"Progress: {current_step}/{len(steps)} steps complete. "
                     f"Failures: {len(failed_steps)}",
                 }
             ]
@@ -312,7 +341,7 @@ def should_continue(
         plan = state.get("plan")
         current_step = state.get("current_step", 0)
 
-        if plan and current_step < len(plan.steps):
+        if plan and current_step < len(_get_plan_steps(plan)):
             return "execute"
 
         return "respond"

langgraph_samples/reflection/graph.py

@@ -196,7 +196,13 @@ def should_revise(state: ReflectionState) -> Literal["revise", "finalize"]:
         latest_critique = critiques[-1]
 
         # Finalize if satisfactory or max iterations reached
-        if latest_critique.is_satisfactory or iteration >= max_iterations:
+        # Note: After Temporal serialization, Critique objects become dicts
+        is_satisfactory = (
+            latest_critique.get("is_satisfactory", False)
+            if isinstance(latest_critique, dict)
+            else latest_critique.is_satisfactory
+        )
+        if is_satisfactory or iteration >= max_iterations:
             return "finalize"
 
         return "revise"
@@ -217,10 +223,20 @@ def revise(state: ReflectionState) -> dict[str, Any]:
         latest_critique = critiques[-1]
 
         # Format the feedback
+        # Note: After Temporal serialization, Critique objects become dicts
+        if isinstance(latest_critique, dict):
+            strengths = latest_critique.get("strengths", [])
+            weaknesses = latest_critique.get("weaknesses", [])
+            suggestions = latest_critique.get("suggestions", [])
+        else:
+            strengths = latest_critique.strengths
+            weaknesses = latest_critique.weaknesses
+            suggestions = latest_critique.suggestions
+
         feedback = f"""
-Strengths: {', '.join(latest_critique.strengths)}
-Weaknesses: {', '.join(latest_critique.weaknesses)}
-Suggestions: {', '.join(latest_critique.suggestions)}
+Strengths: {', '.join(strengths)}
+Weaknesses: {', '.join(weaknesses)}
+Suggestions: {', '.join(suggestions)}
 """
 
         revise_prompt = ChatPromptTemplate.from_messages(
@@ -264,7 +280,16 @@ def finalize(state: ReflectionState) -> dict[str, Any]:
         iteration = state.get("iteration", 1)
 
         # Get final score
-        final_score = critiques[-1].quality_score if critiques else 0
+        # Note: After Temporal serialization, Critique objects become dicts
+        if critiques:
+            last_critique = critiques[-1]
+            final_score = (
+                last_critique.get("quality_score", 0)
+                if isinstance(last_critique, dict)
+                else last_critique.quality_score
+            )
+        else:
+            final_score = 0
 
         summary = f"""
 Content finalized after {iteration} iteration(s).

tests/langgraph_samples/agentic_rag_test.py

@@ -0,0 +1,40 @@
+"""Tests for the agentic_rag LangGraph sample."""
+
+import uuid
+
+from temporalio.client import Client
+from temporalio.contrib.langgraph import LangGraphPlugin
+from temporalio.worker import Worker
+
+from langgraph_samples.agentic_rag.graph import build_agentic_rag_graph
+from langgraph_samples.agentic_rag.workflow import AgenticRAGWorkflow
+
+from .conftest import requires_openai
+
+
+@requires_openai
+async def test_agentic_rag_workflow(client: Client) -> None:
+    """Test agentic RAG workflow with a knowledge base query.
+
+    This test requires OPENAI_API_KEY to be set.
+    """
+    task_queue = f"agentic-rag-test-{uuid.uuid4()}"
+
+    plugin = LangGraphPlugin(graphs={"agentic_rag": build_agentic_rag_graph})
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[AgenticRAGWorkflow],
+        plugins=[plugin],
+    ):
+        result = await client.execute_workflow(
+            AgenticRAGWorkflow.run,
+            "What are AI agents?",
+            id=f"agentic-rag-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+
+        # The result should contain messages with retrieved context
+        assert "messages" in result
+        assert len(result["messages"]) > 0

tests/langgraph_samples/approval_workflow_test.py

@@ -0,0 +1,159 @@
+"""Tests for the approval_workflow LangGraph sample."""
+
+import uuid
+
+from temporalio.client import Client, WorkflowHandle
+from temporalio.contrib.langgraph import LangGraphPlugin
+from temporalio.worker import Worker
+
+from langgraph_samples.approval_workflow.activities import notify_approver
+from langgraph_samples.approval_workflow.graph import build_approval_graph
+from langgraph_samples.approval_workflow.workflow import (
+    ApprovalRequest,
+    ApprovalWorkflow,
+)
+
+
+async def test_approval_workflow_approved(client: Client) -> None:
+    """Test approval workflow when request is approved."""
+    task_queue = f"approval-test-{uuid.uuid4()}"
+
+    plugin = LangGraphPlugin(graphs={"approval_workflow": build_approval_graph})
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[ApprovalWorkflow],
+        activities=[notify_approver],
+        plugins=[plugin],
+    ):
+        # Start the workflow
+        handle: WorkflowHandle[ApprovalWorkflow, dict] = await client.start_workflow(
+            ApprovalWorkflow.run,
+            ApprovalRequest(request_type="expense", amount=500.0),
+            id=f"approval-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+
+        # Wait for the workflow to reach the approval point
+        import asyncio
+
+        for _ in range(20):
+            status = await handle.query(ApprovalWorkflow.get_status)
+            if status == "waiting_for_approval":
+                break
+            await asyncio.sleep(0.1)
+
+        assert status == "waiting_for_approval"
+
+        # Query the pending approval
+        pending = await handle.query(ApprovalWorkflow.get_pending_approval)
+        assert pending is not None
+        assert pending["amount"] == 500.0
+        assert pending["risk_level"] == "medium"
+
+        # Send approval signal
+        await handle.signal(
+            ApprovalWorkflow.provide_approval,
+            {"approved": True, "reason": "Looks good", "approver": "manager"},
+        )
+
+        # Wait for result
+        result = await handle.result()
+
+        assert result["approved"] is True
+        assert result["executed"] is True
+        assert "Successfully processed" in result["result"]
+        assert "manager" in result["result"]
+
+
+async def test_approval_workflow_rejected(client: Client) -> None:
+    """Test approval workflow when request is rejected."""
+    task_queue = f"approval-test-{uuid.uuid4()}"
+
+    plugin = LangGraphPlugin(graphs={"approval_workflow": build_approval_graph})
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[ApprovalWorkflow],
+        activities=[notify_approver],
+        plugins=[plugin],
+    ):
+        handle: WorkflowHandle[ApprovalWorkflow, dict] = await client.start_workflow(
+            ApprovalWorkflow.run,
+            ApprovalRequest(request_type="purchase", amount=5000.0),
+            id=f"approval-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+
+        # Wait for approval state
+        import asyncio
+
+        for _ in range(20):
+            status = await handle.query(ApprovalWorkflow.get_status)
+            if status == "waiting_for_approval":
+                break
+            await asyncio.sleep(0.1)
+
+        # Verify high risk level for large amount
+        pending = await handle.query(ApprovalWorkflow.get_pending_approval)
+        assert pending is not None
+        assert pending["risk_level"] == "high"
+
+        # Reject the request
+        await handle.signal(
+            ApprovalWorkflow.provide_approval,
+            {"approved": False, "reason": "Budget exceeded", "approver": "cfo"},
+        )
+
+        result = await handle.result()
+
+        assert result["approved"] is False
+        assert result["executed"] is False
+        assert "rejected" in result["result"]
+        assert "cfo" in result["result"]
+
+
+async def test_approval_workflow_low_risk(client: Client) -> None:
+    """Test approval workflow with low risk amount."""
+    task_queue = f"approval-test-{uuid.uuid4()}"
+
+    plugin = LangGraphPlugin(graphs={"approval_workflow": build_approval_graph})
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[ApprovalWorkflow],
+        activities=[notify_approver],
+        plugins=[plugin],
+    ):
+        handle: WorkflowHandle[ApprovalWorkflow, dict] = await client.start_workflow(
+            ApprovalWorkflow.run,
+            ApprovalRequest(request_type="supplies", amount=25.0),
+            id=f"approval-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+
+        # Wait for approval state
+        import asyncio
+
+        for _ in range(20):
+            status = await handle.query(ApprovalWorkflow.get_status)
+            if status == "waiting_for_approval":
+                break
+            await asyncio.sleep(0.1)
+
+        # Verify low risk level
+        pending = await handle.query(ApprovalWorkflow.get_pending_approval)
+        assert pending is not None
+        assert pending["risk_level"] == "low"
+
+        # Approve
+        await handle.signal(
+            ApprovalWorkflow.provide_approval,
+            {"approved": True, "reason": "Auto-approved", "approver": "system"},
+        )
+
+        result = await handle.result()
+        assert result["approved"] is True

tests/langgraph_samples/conftest.py

@@ -0,0 +1,27 @@
+"""Shared test fixtures for LangGraph samples."""
+
+import os
+
+# Disable LangSmith tracing for tests to avoid rate limit issues
+os.environ["LANGCHAIN_TRACING_V2"] = "false"
+
+import pytest
+from temporalio.contrib.langgraph._graph_registry import get_global_registry
+
+
+@pytest.fixture(autouse=True)
+def clear_registry() -> None:
+    """Clear the global graph registry before each test."""
+    get_global_registry().clear()
+
+
+def has_openai_api_key() -> bool:
+    """Check if OpenAI API key is available."""
+    return bool(os.environ.get("OPENAI_API_KEY"))
+
+
+# Skip marker for tests that require OpenAI API
+requires_openai = pytest.mark.skipif(
+    not has_openai_api_key(),
+    reason="OPENAI_API_KEY environment variable not set",
+)