feat:Make VertexAiSessionService true async.

PiperOrigin-RevId: 762547133
This commit is contained in:
Shangjie Chen
2025-05-23 13:30:57 -07:00
committed by Copybara-Service
parent 79681e3513
commit d212e50c10
2 changed files with 47 additions and 17 deletions

View File

@@ -15,6 +15,7 @@
import re
import this
from typing import Any
from unittest import mock
from dateutil.parser import isoparse
from google.adk.events import Event
@@ -123,7 +124,9 @@ class MockApiClient:
this.session_dict: dict[str, Any] = {}
this.event_dict: dict[str, list[Any]] = {}
def request(self, http_method: str, path: str, request_dict: dict[str, Any]):
async def async_request(
self, http_method: str, path: str, request_dict: dict[str, Any]
):
"""Mocks the API Client request method."""
if http_method == 'GET':
if re.match(SESSION_REGEX, path):
@@ -194,22 +197,31 @@ class MockApiClient:
def mock_vertex_ai_session_service():
"""Creates a mock Vertex AI Session service for testing."""
service = VertexAiSessionService(
return VertexAiSessionService(
project='test-project', location='test-location'
)
service.api_client = MockApiClient()
service.api_client.session_dict = {
@pytest.fixture
def mock_get_api_client():
api_client = MockApiClient()
api_client.session_dict = {
'1': MOCK_SESSION_JSON_1,
'2': MOCK_SESSION_JSON_2,
'3': MOCK_SESSION_JSON_3,
}
service.api_client.event_dict = {
api_client.event_dict = {
'1': MOCK_EVENT_JSON,
}
return service
with mock.patch(
"google.adk.sessions.vertex_ai_session_service._get_api_client",
return_value=api_client,
):
yield
@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_get_empty_session():
session_service = mock_vertex_ai_session_service()
with pytest.raises(ValueError) as excinfo:
@@ -220,6 +232,7 @@ async def test_get_empty_session():
@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_get_and_delete_session():
session_service = mock_vertex_ai_session_service()
@@ -241,6 +254,7 @@ async def test_get_and_delete_session():
@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_list_sessions():
session_service = mock_vertex_ai_session_service()
sessions = await session_service.list_sessions(app_name='123', user_id='user')
@@ -250,6 +264,7 @@ async def test_list_sessions():
@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_create_session():
session_service = mock_vertex_ai_session_service()
@@ -269,6 +284,7 @@ async def test_create_session():
@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_create_session_with_custom_session_id():
session_service = mock_vertex_ai_session_service()