Improving tests and handling one failure mode in memory

Author: ScottMansfield

src/google/adk/integrations/firestore/firestore_memory_service.py

@@ -169,11 +169,14 @@ async def search_memory(
         self._search_by_keyword(app_name, user_id, keyword)
         for keyword in keywords
     ]
-    results = await asyncio.gather(*tasks)
+    results = await asyncio.gather(*tasks, return_exceptions=True)
 
     seen = set()
     memories = []
     for result_list in results:
+      if isinstance(result_list, Exception):
+        logger.warning(f"Memory keyword search partial failure: {result_list}")
+        continue
       for entry in result_list:
         content_text = ""
         if entry.content and entry.content.parts:

tests/unittests/integrations/firestore/test_firestore_memory_service.py

@@ -201,6 +201,41 @@ async def test_search_memory_only_stop_words(mock_firestore_client):
   mock_firestore_client.collection.assert_not_called()
 
 
+@pytest.mark.asyncio
+async def test_search_memory_partial_failures(mock_firestore_client, caplog):
+  service = FirestoreMemoryService(client=mock_firestore_client)
+  app_name = "test_app"
+  user_id = "test_user"
+  query = "fox quick"
+
+  coll_ref = mock_firestore_client.collection.return_value.where.return_value.where.return_value.where.return_value
+  
+  doc_snapshot = mock.MagicMock()
+  doc_snapshot.to_dict.return_value = {
+      "content": {"parts": [{"text": "quick response"}]},
+      "author": "user",
+      "timestamp": 1234567890.0
+  }
+
+  call_count = 0
+  async def mock_get():
+    nonlocal call_count
+    call_count += 1
+    if call_count == 1:
+       raise ValueError("Mock generic network failure standalone")
+    return [doc_snapshot]
+
+  coll_ref.get = mock.AsyncMock(side_effect=mock_get)
+
+  response = await service.search_memory(
+      app_name=app_name, user_id=user_id, query=query
+  )
+
+  assert len(response.memories) == 1
+  assert response.memories[0].author == "user"
+  assert "Memory keyword search partial failure" in caplog.text
+
+
 def test_init_default_client():
   with mock.patch("google.cloud.firestore.AsyncClient") as mock_client_class:
     mock_instance = mock.MagicMock()

tests/unittests/integrations/firestore/test_firestore_session_service.py

@@ -427,6 +427,63 @@ def collection_side_effect(name):
   )
 
 
+@pytest.mark.asyncio
+async def test_list_sessions_filters_other_apps(mock_firestore_client):
+  service = FirestoreSessionService(client=mock_firestore_client)
+  app_name = "test_app"
+
+  session_doc = mock.MagicMock()
+  session_doc.to_dict.return_value = {
+      "id": "session1",
+      "appName": app_name,
+      "userId": "user1",
+      "state": {"session_key": "session_val"},
+  }
+
+  mock_firestore_client.collection_group.return_value.where.return_value.get = (
+      mock.AsyncMock(return_value=[session_doc])
+  )
+
+  app_state_coll = mock.MagicMock()
+  user_state_coll = mock.MagicMock()
+
+  def collection_side_effect(name):
+    if name == service.app_state_collection:
+      return app_state_coll
+    elif name == service.user_state_collection:
+      return user_state_coll
+    return mock.MagicMock()
+
+  mock_firestore_client.collection.side_effect = collection_side_effect
+
+  app_doc = mock.MagicMock()
+  app_doc.exists = True
+  app_doc.to_dict.return_value = {"app_key": "app_val"}
+  app_doc_ref = mock.MagicMock()
+  app_state_coll.document.return_value = app_doc_ref
+  app_doc_ref.get = mock.AsyncMock(return_value=app_doc)
+
+  user_doc = mock.MagicMock()
+  user_doc.id = "user1"
+  user_doc.to_dict.return_value = {"user_key": "user_val"}
+  user_app_doc = mock.MagicMock()
+  user_state_coll.document.return_value = user_app_doc
+  users_coll = mock.MagicMock()
+  user_app_doc.collection.return_value = users_coll
+  users_coll.get = mock.AsyncMock(return_value=[user_doc])
+
+  response = await service.list_sessions(app_name=app_name)
+
+  assert len(response.sessions) == 1
+  assert response.sessions[0].id == "session1"
+  assert response.sessions[0].app_name == app_name
+
+  mock_firestore_client.collection_group.assert_called_once_with("sessions")
+  mock_firestore_client.collection_group.return_value.where.assert_called_once_with(
+      "appName", "==", app_name
+  )
+
+
 @pytest.mark.asyncio
 async def test_create_session_already_exists(mock_firestore_client):
   service = FirestoreSessionService(client=mock_firestore_client)