Robust Testing for Long-Running Agent Workflows

.
├── my_adk_agent/
│   ├── __init__.py
│   ├── agent.py               # Our simplified ADK agent model
│   └── persistence.py         # State serialization and Firestore interaction
├── tests/
│   ├── __init__.py
│   ├── test_persistence_units.py  # Unit tests for serialization
│   ├── test_persistence_integration.py # Integration tests for Firestore
│   └── test_e2e_resume.py         # E2E tests for pause/resume workflow
├── requirements.txt
└── main.py

pip install pytest pytest-mock google-cloud-firestore

# tests/test_persistence_units.py
import pytest
from my_adk_agent.persistence import serialize_state, deserialize_state

def test_state_serialization_deserialization():
    """
    Verifies that agent state can be serialized to JSON and deserialized back
    without loss of data or type changes.
    """
    initial_state = {
        "user_id": "test_user_123",
        "conversation_history": [
            {"role": "user", "content": "Hello"},
            {"role": "agent", "content": "Hi there!"}
        ],
        "current_workflow_step": "awaiting_confirmation",
        "context_data": {"order_id": "ORD-001", "item_count": 2},
        "is_active": True,
        "last_updated_timestamp": 1700000000.0, # Example timestamp
    }

    # 📌 Key Idea: Test round-trip fidelity.
    serialized_data = serialize_state(initial_state)
    assert isinstance(serialized_data, str), "Serialized data should be a string."
    assert 'test_user_123' in serialized_data, "User ID should be present in serialized data."
    assert 'awaiting_confirmation' in serialized_data, "Workflow step should be present."

    deserialized_state = deserialize_state(serialized_data)

    # Assert that the deserialized state matches the original
    assert deserialized_state == initial_state, "Deserialized state must match initial state."
    # Also check specific types for robustness, especially for numbers/booleans
    assert isinstance(deserialized_state['is_active'], bool), "Boolean type should be preserved."
    assert isinstance(deserialized_state['last_updated_timestamp'], float), "Float type should be preserved."
    assert deserialized_state['conversation_history'][0]['role'] == 'user', "Nested list/dict data should be correct."

    print(f"\nInitial state: {initial_state}")
    print(f"Serialized data: {serialized_data}")
    print(f"Deserialized state: {deserialized_state}")

def test_empty_state_serialization():
    """Tests serialization/deserialization of an empty state dictionary."""
    empty_state = {}
    serialized = serialize_state(empty_state)
    deserialized = deserialize_state(serialized)
    assert deserialized == empty_state, "Empty state should serialize and deserialize correctly."

def test_invalid_json_deserialization():
    """Tests deserialization with invalid JSON input to ensure error handling."""
    # ⚠️ What can go wrong: Malformed JSON can crash your agent.
    with pytest.raises(ValueError, match="Invalid JSON format"):
        deserialize_state("this is not valid json")

    with pytest.raises(ValueError, match="Invalid JSON format"):
        deserialize_state("{'key': 'value'") # Malformed JSON with single quotes

gcloud components install cloud-firestore-emulator
gcloud emulators firestore start --host-port=localhost:8080

# my_adk_agent/persistence.py (Illustrative snippet, adapt to your implementation)
import json
from google.cloud import firestore

class StateManager:
    """Abstract base class for state management."""
    def save_agent_state(self, agent_id: str, state: dict):
        raise NotImplementedError
    def load_agent_state(self, agent_id: str) -> dict | None:
        raise NotImplementedError

class FirestoreStateManager(StateManager):
    """Manages agent state persistence using Google Cloud Firestore."""
    def __init__(self, client: firestore.Client, collection_name: str = "agent_states"):
        self.db = client
        self.collection_ref = self.db.collection(collection_name)

    def save_agent_state(self, agent_id: str, state: dict):
        """Saves the agent's state to Firestore."""
        doc_ref = self.collection_ref.document(agent_id)
        # Firestore automatically handles Python dicts, converting them to JSON-like documents.
        # We assume `state` is already a dictionary of basic types or objects Firestore can handle.
        doc_ref.set(state)
        # ⚡ Quick Note: For complex objects, you might need pre-serialization here.
        # Our `serialize_state` is more for general JSON string output, not direct Firestore dicts.
        # Firestore's client library handles basic Python types (dicts, lists, int, str, bool, float) well.

    def load_agent_state(self, agent_id: str) -> dict | None:
        """Loads the agent's state from Firestore."""
        doc_ref = self.collection_ref.document(agent_id)
        doc = doc_ref.get()
        if doc.exists:
            return doc.to_dict()
        return None

def serialize_state(state: dict) -> str:
    """Serializes a dictionary state to a JSON string."""
    try:
        return json.dumps(state)
    except TypeError as e:
        raise ValueError(f"Failed to serialize state to JSON: {e}")

def deserialize_state(serialized_state: str) -> dict:
    """Deserializes a JSON string back to a dictionary state."""
    try:
        return json.loads(serialized_state)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON format for state: {e}")

# tests/test_persistence_integration.py
import pytest
import os
from google.cloud import firestore
from my_adk_agent.persistence import FirestoreStateManager

# IMPORTANT: Configure Firestore emulator for tests
# This environment variable tells the client library to connect to the emulator.
os.environ["FIRESTORE_EMULATOR_HOST"] = "localhost:8080"
os.environ["GCLOUD_PROJECT"] = "test-project-id" # Use a dummy project ID for emulator

@pytest.fixture(scope="module")
def firestore_client_fixture():
    """Provides a Firestore client connected to the emulator for the test module."""
    client = firestore.Client()
    yield client
    # Clean up any data created by tests in this module
    # 🔥 Optimization / Pro tip: For more granular cleanup, delete per test or use a unique collection name per test run.
    # For a simple module-scoped cleanup, deleting the collection is sufficient.
    print("\nCleaning up Firestore emulator data...")
    docs = client.collection("agent_states").stream()
    for doc in docs:
        doc.reference.delete()
    print("Firestore emulator data cleaned.")


def test_save_and_load_agent_state_integration(firestore_client_fixture):
    """
    Tests saving and loading agent state to/from Firestore using the emulator.
    This validates the interaction between our StateManager and Firestore.
    """
    manager = FirestoreStateManager(firestore_client_fixture, collection_name="agent_states")
    test_agent_id = "agent_test_id_001"
    initial_state = {
        "user_id": "test_user_001",
        "conversation_history": [{"role": "user", "content": "Hello Firestore"}],
        "workflow_status": "pending_db_check",
        "nested_data": {"sub_key": 123, "list_items": ["a", "b"]}
    }

    # 1. Save state
    manager.save_agent_state(test_agent_id, initial_state)

    # 2. Load state
    loaded_state = manager.load_agent_state(test_agent_id)

    # 3. Assertions
    assert loaded_state is not None, "Loaded state should not be None."
    assert loaded_state == initial_state, "Loaded state must exactly match initial state."
    assert loaded_state["user_id"] == "test_user_001"
    assert loaded_state["conversation_history"][0]["content"] == "Hello Firestore"
    assert loaded_state["nested_data"]["sub_key"] == 123

    # Test updating state: crucial for long-running agents
    updated_state = initial_state.copy()
    updated_state["workflow_status"] = "completed_db_check"
    updated_state["new_field"] = "added_during_update"
    manager.save_agent_state(test_agent_id, updated_state)
    reloaded_updated_state = manager.load_agent_state(test_agent_id)

    assert reloaded_updated_state == updated_state, "Updated state must be correctly saved and loaded."
    assert reloaded_updated_state["new_field"] == "added_during_update"

def test_load_non_existent_state(firestore_client_fixture):
    """
    Tests loading a state for an agent ID that does not exist.
    It should gracefully return None, indicating no state found.
    """
    manager = FirestoreStateManager(firestore_client_fixture, collection_name="agent_states")
    non_existent_id = "non_existent_agent_id_123"
    loaded_state = manager.load_agent_state(non_existent_id)
    assert loaded_state is None, "Loading non-existent state should return None."

# my_adk_agent/agent.py (Illustrative snippet, adapt to your ADK agent structure)
from .persistence import StateManager, FirestoreStateManager # Assuming FirestoreStateManager is used
# from adk.agent import Agent as AdkAgent # Placeholder for ADK Agent class, not directly used in this simplified model
# from adk.message import Message # Placeholder for ADK Message class

class MyADKAgent:
    """
    A simplified agent model for demonstrating state management and pause/resume.
    In a real ADK setup, state management would integrate with ADK's lifecycle hooks.
    This model abstracts the ADK specifics to focus on state persistence testing.
    """
    def __init__(self, agent_id: str, state_manager: StateManager):
        self.agent_id = agent_id
        self.state_manager = state_manager
        # Load existing state or initialize new state
        self.state = self.state_manager.load_agent_state(self.agent_id) or self._initial_state()
        print(f"Agent {self.agent_id} initialized. Current state: {self.state.get('current_workflow_step')}")


    def _initial_state(self) -> dict:
        """Returns the default initial state for a new agent session."""
        return {
            "conversation_history": [],
            "current_workflow_step": "start",
            "context_data": {}
        }

    def _save_state(self):
        """Saves the current internal state of the agent using the state manager."""
        self.state_manager.save_agent_state(self.agent_id, self.state)

    def handle_message(self, user_message: str) -> str:
        """
        Simulates agent processing of a user message and updates its internal state.
        Saves state after each interaction.
        """
        self.state["conversation_history"].append({"role": "user", "content": user_message})

        response = ""
        current_step = self.state.get("current_workflow_step")

        # 🧠 Important: This simplified logic demonstrates state transitions.
        # A real ADK agent would use LLM calls, tool execution, etc.
        if current_step == "start":
            response = "Welcome! What is your name?"
            self.state["current_workflow_step"] = "awaiting_name"
        elif current_step == "awaiting_name":
            self.state["context_data"]["user_name"] = user_message
            response = f"Nice to meet you, {user_message}. What task can I help you with?"
            self.state["current_workflow_step"] = "awaiting_task"
        elif current_step == "awaiting_task":
            response = f"Understood. I will help with '{user_message}'. Is that correct? (yes/no)"
            self.state["context_data"]["proposed_task"] = user_message
            self.state["current_workflow_step"] = "confirm_task"
        elif current_step == "confirm_task":
            if user_message.lower() == "yes":
                response = "Great! Starting your task now."
                self.state["current_workflow_step"] = "task_in_progress"
            else:
                response = "Okay, what task should I help with instead?"
                self.state["current_workflow_step"] = "awaiting_task"
        elif current_step == "task_in_progress":
            response = "I'm currently working on your task. What's next?"
        else:
            response = "I'm not sure how to proceed. Let's restart. What is your name?"
            self.state["current_workflow_step"] = "awaiting_name"

        self.state["conversation_history"].append({"role": "agent", "content": response})
        self._save_state() # Save state after every interaction to ensure persistence
        return response

    def get_current_workflow_step(self) -> str:
        """Returns the current step in the agent's workflow."""
        return self.state.get("current_workflow_step", "unknown")

    def get_context_data(self) -> dict:
        """Returns the current context data maintained by the agent."""
        return self.state.get("context_data", {})

# tests/test_e2e_resume.py
import pytest
import os
from google.cloud import firestore
from my_adk_agent.persistence import FirestoreStateManager
from my_adk_agent.agent import MyADKAgent

# Ensure Firestore emulator is configured for tests
os.environ["FIRESTORE_EMULATOR_HOST"] = "localhost:8080"
os.environ["GCLOUD_PROJECT"] = "test-project-id"

@pytest.fixture(scope="module")
def firestore_client_e2e_fixture():
    """Provides a Firestore client connected to the emulator for E2E tests."""
    client = firestore.Client()
    yield client
    # Clean up specific collection used by E2E tests
    print("\nCleaning up E2E Firestore emulator data...")
    docs = client.collection("e2e_agent_states").stream()
    for doc in docs:
        doc.reference.delete()
    print("E2E Firestore emulator data cleaned.")

def test_agent_pause_resume_workflow(firestore_client_e2e_fixture):
    """
    Tests an end-to-end scenario where an agent workflow is paused and then resumed
    from the correct state and context. This simulates an agent being restarted.
    """
    test_agent_id = "e2e_agent_id_001"
    state_manager = FirestoreStateManager(firestore_client_e2e_fixture, collection_name="e2e_agent_states")

    # --- Part 1: Initial interaction and state saving (Simulate Agent A) ---
    print("\n--- Agent A: Initial conversation ---")
    agent_a = MyADKAgent(test_agent_id, state_manager)
    # Agent A starts, loads state (will be initial state as it's the first run)

    # First interaction
    response1 = agent_a.handle_message("My name is Alice")
    print(f"Agent A Response 1: {response1}")
    assert "Nice to meet you, Alice" in response1, "Agent A should greet Alice."
    assert agent_a.get_current_workflow_step() == "awaiting_task", "Agent A should move to awaiting_task step."
    assert agent_a.get_context_data().get("user_name") == "Alice", "Agent A should store user name."

    # Second interaction
    response2 = agent_a.handle_message("I need help with booking a flight")
    print(f"Agent A Response 2: {response2}")
    assert "booking a flight" in response2, "Agent A should confirm task."
    assert agent_a.get_current_workflow_step() == "confirm_task", "Agent A should move to confirm_task step."
    assert agent_a.get_context_data().get("proposed_task") == "booking a flight", "Agent A should store proposed task."

    # At this point, agent_a's state is saved to Firestore by `handle_message`
    print(f"Agent A current state: {agent_a.get_current_workflow_step()}, context: {agent_a.get_context_data()}")

    # --- Part 2: Simulate agent shutdown and restart (Simulate Agent B) ---
    print("\n--- Agent B: Resuming conversation ---")
    # Simulate a new instance of the agent, loading state from persistence
    agent_b = MyADKAgent(test_agent_id, state_manager)
    # Agent B initializes, and crucially, loads the state saved by Agent A

    # Verify agent_b has correctly loaded the state from agent_a
    assert agent_b.get_current_workflow_step() == "confirm_task", "Agent B should resume at 'confirm_task' step."
    assert agent_b.get_context_data().get("user_name") == "Alice", "Agent B should have Alice's name in context."
    assert agent_b.get_context_data().get("proposed_task") == "booking a flight", "Agent B should have proposed task in context."
    print(f"Agent B loaded workflow step: {agent_b.get_current_workflow_step()}")
    print(f"Agent B loaded context data: {agent_b.get_context_data()}")

    # Continue the conversation from where Agent A left off
    response3 = agent_b.handle_message("yes")
    print(f"Agent B Response 3: {response3}")
    assert "Starting your task now" in response3, "Agent B should proceed to task start confirmation."
    assert agent_b.get_current_workflow_step() == "task_in_progress", "Agent B should move to 'task_in_progress' step."
    assert agent_b.get_context_data().get("proposed_task") == "booking a flight", "Agent B should retain proposed task after confirmation."

    print("\nE2E Pause/Resume Test Passed!")

pytest

============================= test session starts ==============================
platform linux -- Python 3.12.x, pytest-X.X.X, pluggy-X.X.X
rootdir: /path/to/your/project
plugins: mock-X.X.X
collected 6 items

tests/test_persistence_units.py::test_state_serialization_deserialization PASSED [ 16%]
tests/test_persistence_units.py::test_empty_state_serialization PASSED [ 33%]
tests/test_persistence_units.py::test_invalid_json_deserialization PASSED [ 50%]
tests/test_persistence_integration.py::test_save_and_load_agent_state_integration PASSED [ 66%]
tests/test_persistence_integration.py::test_load_non_existent_state PASSED [ 83%]
tests/test_e2e_resume.py::test_agent_pause_resume_workflow PASSED [100%]

============================== 6 passed in X.XXs ===============================