	In Browser
	StumbleUpon
	del.icio.us
	Google
	Google Buzz
	reddit
	LinkedIn

	Facebook
	Twitter
	Linkedin
	E-Mail

Agentic AI > ADK Agent Testing > Testing ADK Agents with Firestore Session Persistence

Testing ADK Agents with Firestore Session Persistence

Author: Venkata Sudhakar

ShopMax India uses Firestore to persist ADK agent session state across service restarts and horizontal scaling - ensuring a customer in Chennai who starts a return request on one pod can seamlessly continue on another. Testing Firestore session persistence verifies that state written during one agent invocation is correctly read back in the next, and that Firestore-specific behaviors like document TTL and partial updates work as expected in the agent context.

To test Firestore-backed ADK sessions without a live database, use a mock Firestore client that stores documents in memory and mirrors the Firestore API: set, get, update, and delete by document path. Assert that the session service writes to the correct collection and document path, that subsequent reads return the persisted state, and that TTL fields are written correctly. This approach runs fast in CI without Firestore credentials.

The example below implements a mock Firestore client and a Firestore-backed session service for ShopMax India. Tests verify write path correctness, state persistence across reads, partial updates via merge, and TTL field presence on every written document.

import pytest
import time

class MockFirestore:
    def __init__(self): self._docs = {}

def collection(self, name): return MockCollection(self._docs, name)

class MockCollection:
    def __init__(self, store, name): self._store = store; self._name = name

def document(self, doc_id): return MockDocument(self._store, self._name + "/" + doc_id)

class MockDocument:
    def __init__(self, store, path): self._store = store; self._path = path

def set(self, data, merge=False):
        if merge and self._path in self._store:
            self._store[self._path].update(data)
        else:
            self._store[self._path] = dict(data)

def get(self):
        data = self._store.get(self._path)
        return MockDocSnapshot(data)

class MockDocSnapshot:
    def __init__(self, data): self._data = data
    def to_dict(self): return dict(self._data) if self._data else None

class FirestoreSessionService:
    def __init__(self, db, collection="agent_sessions"):
        self._db = db; self._col = collection

def save(self, sid, state):
        data = dict(state)
        data["_ttl"] = int(time.time()) + 1800
        self._db.collection(self._col).document(sid).set(data)

def load(self, sid):
        snap = self._db.collection(self._col).document(sid).get()
        return snap.to_dict()

def update(self, sid, partial):
        self._db.collection(self._col).document(sid).set(partial, merge=True)

def test_save_and_load_session():
    svc = FirestoreSessionService(MockFirestore())
    svc.save("SID-FS-001", {"customer": "Arjun", "city": "Bangalore", "order": "ORD-2201"})
    state = svc.load("SID-FS-001")
    assert state["customer"] == "Arjun"
    assert state["city"] == "Bangalore"

def test_partial_update_merges_state():
    svc = FirestoreSessionService(MockFirestore())
    svc.save("SID-FS-002", {"order": "ORD-2202", "status": "Pending"})
    svc.update("SID-FS-002", {"status": "Confirmed"})
    state = svc.load("SID-FS-002")
    assert state["order"] == "ORD-2202"
    assert state["status"] == "Confirmed"

def test_ttl_field_written_on_save():
    svc = FirestoreSessionService(MockFirestore())
    svc.save("SID-FS-003", {"customer": "Meena", "city": "Chennai"})
    state = svc.load("SID-FS-003")
    assert "_ttl" in state
    assert state["_ttl"] > int(time.time())

It gives the following output,

... (3 passed in 0.01s)

In production, ShopMax India should use Firestore TTL policies to automatically delete expired sessions without a manual cleanup job. Index the session collection on customer ID so support agents can look up all active sessions for a given customer. Use Firestore transactions when updating session state concurrently to prevent write conflicts during the Diwali sale when thousands of customers update their carts simultaneously.

Send your comments, suggestions or queries regarding this site to [email protected].