refactor: update list_sessions to use firestore get_all method

Author: Dumeng

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

@@ -386,10 +386,8 @@ async def list_sessions(
             .document(app_name)
             .collection("users")
         )
-        user_docs = await asyncio.gather(
-            *[users_coll.document(uid).get() for uid in unique_user_ids]
-        )
-        for u_doc in user_docs:
+        refs = [users_coll.document(uid) for uid in sorted(unique_user_ids)]
+        async for u_doc in self.client.get_all(refs):
           if u_doc.exists:
             user_states_map[u_doc.id] = u_doc.to_dict()
 

tests/unittests/integrations/firestore/test_firestore_session_service.py

@@ -482,7 +482,11 @@ def collection_side_effect(name):
   user_app_doc.collection.return_value = users_coll
   user_doc_ref = mock.MagicMock()
   users_coll.document.return_value = user_doc_ref
-  user_doc_ref.get = mock.AsyncMock(return_value=user_doc)
+
+  async def mock_get_all(refs):
+    yield user_doc
+
+  mock_firestore_client.get_all = mock_get_all
 
   response = await service.list_sessions(app_name=app_name)
 
@@ -544,7 +548,11 @@ def collection_side_effect(name):
   user_app_doc.collection.return_value = users_coll
   user_doc_ref = mock.MagicMock()
   users_coll.document.return_value = user_doc_ref
-  user_doc_ref.get = mock.AsyncMock(return_value=user_doc)
+
+  async def mock_get_all(refs):
+    yield user_doc
+
+  mock_firestore_client.get_all = mock_get_all
 
   response = await service.list_sessions(app_name=app_name)