Fix async mocks

Author: ScottMansfield

tests/unittests/memory/test_firestore_memory_service.py

@@ -24,15 +24,18 @@
 
 @pytest.fixture
 def mock_firestore_client():
-  client = mock.AsyncMock()
-  collection_ref = mock.AsyncMock()
+  client = mock.MagicMock()
+  collection_ref = mock.MagicMock()
   client.collection_group.return_value = collection_ref
+
+  # where() should return self (collection_ref) to allow chaining
   collection_ref.where.return_value = collection_ref
 
   # Mock get() for documents
-  doc_snapshot = mock.AsyncMock()
+  doc_snapshot = mock.MagicMock()
   doc_snapshot.to_dict.return_value = {}
-  collection_ref.get.return_value = [doc_snapshot]
+
+  collection_ref.get = mock.AsyncMock(return_value=[doc_snapshot])
 
   return client
 

tests/unittests/sessions/test_firestore_session_service.py

@@ -23,28 +23,36 @@
 
 @pytest.fixture
 def mock_firestore_client():
-  client = mock.AsyncMock()
-  # Mock collection and document references
-  collection_ref = mock.AsyncMock()
-  doc_ref = mock.AsyncMock()
-  subcollection_ref = mock.AsyncMock()
-  subdoc_ref = mock.AsyncMock()
+  client = mock.MagicMock()
+  collection_ref = mock.MagicMock()
+  doc_ref = mock.MagicMock()
+  subcollection_ref = mock.MagicMock()
+  subdoc_ref = mock.MagicMock()
 
   client.collection.return_value = collection_ref
   collection_ref.document.return_value = doc_ref
   doc_ref.collection.return_value = subcollection_ref
   subcollection_ref.document.return_value = subdoc_ref
 
   # Mock get() for documents
-  doc_snapshot = mock.AsyncMock()
+  doc_snapshot = mock.MagicMock()
   doc_snapshot.exists = False
   doc_snapshot.to_dict.return_value = {}
-  doc_ref.get.return_value = doc_snapshot
-  subdoc_ref.get.return_value = doc_snapshot
+
+  doc_ref.get = mock.AsyncMock(return_value=doc_snapshot)
+  subdoc_ref.get = mock.AsyncMock(return_value=doc_snapshot)
+
+  # Mock subcollection get() (for events list in delete_session)
+  subcollection_ref.get = mock.AsyncMock(return_value=[])
 
   # Mock collection group
   client.collection_group.return_value = collection_ref
 
+  # Mock batch
+  batch = mock.MagicMock()
+  client.batch.return_value = batch
+  batch.commit = mock.AsyncMock()
+
   return client
 
 

tests/unittests/test_firestore_database_runner.py

@@ -31,10 +31,15 @@ def mock_agent():
 def test_create_firestore_runner_with_arg(mock_agent, monkeypatch):
   monkeypatch.delenv("ADK_GCS_BUCKET_NAME", raising=False)
 
-  # Mock GcsArtifactService to avoid real client init
-  with mock.patch(
-      "google.adk.firestore_database_runner.GcsArtifactService"
-  ) as mock_gcs:
+  with (
+      mock.patch(
+          "google.adk.firestore_database_runner.FirestoreSessionService"
+      ),
+      mock.patch("google.adk.firestore_database_runner.FirestoreMemoryService"),
+      mock.patch(
+          "google.adk.firestore_database_runner.GcsArtifactService"
+      ) as mock_gcs,
+  ):
     runner = create_firestore_runner(mock_agent, gcs_bucket_name="test_bucket")
 
     assert runner is not None
@@ -44,9 +49,15 @@ def test_create_firestore_runner_with_arg(mock_agent, monkeypatch):
 def test_create_firestore_runner_with_env(mock_agent, monkeypatch):
   monkeypatch.setenv("ADK_GCS_BUCKET_NAME", "env_bucket")
 
-  with mock.patch(
-      "google.adk.firestore_database_runner.GcsArtifactService"
-  ) as mock_gcs:
+  with (
+      mock.patch(
+          "google.adk.firestore_database_runner.FirestoreSessionService"
+      ),
+      mock.patch("google.adk.firestore_database_runner.FirestoreMemoryService"),
+      mock.patch(
+          "google.adk.firestore_database_runner.GcsArtifactService"
+      ) as mock_gcs,
+  ):
     runner = create_firestore_runner(mock_agent)
 
     assert runner is not None