google--adk-python
wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
2055 行
66 KiB
Python
2055 行
66 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from datetime import datetime
|
|
from datetime import timezone
|
|
import enum
|
|
import sqlite3
|
|
from unittest import mock
|
|
|
|
from google.adk.errors.already_exists_error import AlreadyExistsError
|
|
from google.adk.events.event import Event
|
|
from google.adk.events.event_actions import EventActions
|
|
from google.adk.features import FeatureName
|
|
from google.adk.features import override_feature_enabled
|
|
from google.adk.sessions import database_session_service
|
|
from google.adk.sessions.base_session_service import GetSessionConfig
|
|
from google.adk.sessions.database_session_service import DatabaseSessionService
|
|
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
|
from google.adk.sessions.sqlite_session_service import SqliteSessionService
|
|
from google.adk.sessions.vertex_ai_session_service import VertexAiSessionService
|
|
from google.genai import types
|
|
import pytest
|
|
from sqlalchemy import delete
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
|
|
class SessionServiceType(enum.Enum):
|
|
IN_MEMORY = 'IN_MEMORY'
|
|
IN_MEMORY_WITH_LIGHT_COPY_ENABLED = 'IN_MEMORY_WITH_LIGHT_COPY_ENABLED'
|
|
DATABASE = 'DATABASE'
|
|
SQLITE = 'SQLITE'
|
|
|
|
|
|
def get_session_service(
|
|
service_type: SessionServiceType = SessionServiceType.IN_MEMORY,
|
|
tmp_path=None,
|
|
):
|
|
"""Creates a session service for testing."""
|
|
if service_type == SessionServiceType.DATABASE:
|
|
return DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
if service_type == SessionServiceType.SQLITE:
|
|
return SqliteSessionService(str(tmp_path / 'sqlite.db'))
|
|
if service_type == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
|
|
return InMemorySessionService()
|
|
return InMemorySessionService()
|
|
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
SessionServiceType.IN_MEMORY,
|
|
SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED,
|
|
SessionServiceType.DATABASE,
|
|
SessionServiceType.SQLITE,
|
|
]
|
|
)
|
|
async def session_service(request, tmp_path):
|
|
"""Provides a session service and closes database backends on teardown."""
|
|
if request.param == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
|
|
override_feature_enabled(
|
|
FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, True
|
|
)
|
|
service = get_session_service(request.param, tmp_path)
|
|
yield service
|
|
if isinstance(service, DatabaseSessionService):
|
|
await service.close()
|
|
if request.param == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
|
|
override_feature_enabled(
|
|
FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, False
|
|
)
|
|
|
|
|
|
def test_database_session_service_enables_pool_pre_ping_by_default():
|
|
captured_kwargs = {}
|
|
|
|
def fake_create_async_engine(_db_url: str, **kwargs):
|
|
captured_kwargs.update(kwargs)
|
|
fake_engine = mock.Mock()
|
|
fake_engine.dialect.name = 'postgresql'
|
|
fake_engine.sync_engine = mock.Mock()
|
|
return fake_engine
|
|
|
|
with mock.patch.object(
|
|
database_session_service,
|
|
'create_async_engine',
|
|
side_effect=fake_create_async_engine,
|
|
):
|
|
database_session_service.DatabaseSessionService(
|
|
'postgresql+psycopg2://user:pass@localhost:5432/db'
|
|
)
|
|
|
|
assert captured_kwargs.get('pool_pre_ping') is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
'dialect_name', ['sqlite', 'postgresql', 'mysql', 'mariadb']
|
|
)
|
|
def test_database_session_service_uses_naive_datetime_for_dialect(dialect_name):
|
|
"""Verifies dialects that store DATETIME WITHOUT TIME ZONE are treated as naive.
|
|
|
|
SQLite, PostgreSQL, MySQL, and MariaDB all store DATETIME/TIMESTAMP WITHOUT
|
|
TIME ZONE, so create_session must strip tzinfo before storing. Otherwise the
|
|
marker produced by create_session (with +00:00) mismatches the marker read
|
|
back from storage (without +00:00), triggering a false stale-writer error on
|
|
the first append_event after create_session.
|
|
|
|
This exercises the production decision (_uses_naive_datetime) directly rather
|
|
than re-implementing the strip logic, so it actually guards create_session.
|
|
"""
|
|
fake_engine = mock.Mock()
|
|
fake_engine.dialect.name = dialect_name
|
|
fake_engine.sync_engine = mock.Mock()
|
|
|
|
service = DatabaseSessionService(db_engine=fake_engine)
|
|
|
|
assert service._uses_naive_datetime() is True
|
|
|
|
|
|
def test_database_session_service_keeps_timezone_for_spanner():
|
|
"""Spanner's TIMESTAMP is timezone-aware, so tzinfo must be preserved."""
|
|
fake_engine = mock.Mock()
|
|
fake_engine.dialect.name = 'spanner+spanner'
|
|
fake_engine.sync_engine = mock.Mock()
|
|
|
|
service = DatabaseSessionService(db_engine=fake_engine)
|
|
|
|
assert service._uses_naive_datetime() is False
|
|
|
|
|
|
def test_database_session_service_respects_pool_pre_ping_override():
|
|
captured_kwargs = {}
|
|
|
|
def fake_create_async_engine(_db_url: str, **kwargs):
|
|
captured_kwargs.update(kwargs)
|
|
fake_engine = mock.Mock()
|
|
fake_engine.dialect.name = 'postgresql'
|
|
fake_engine.sync_engine = mock.Mock()
|
|
return fake_engine
|
|
|
|
with mock.patch.object(
|
|
database_session_service,
|
|
'create_async_engine',
|
|
side_effect=fake_create_async_engine,
|
|
):
|
|
database_session_service.DatabaseSessionService(
|
|
'postgresql+psycopg2://user:pass@localhost:5432/db',
|
|
pool_pre_ping=False,
|
|
)
|
|
|
|
assert captured_kwargs.get('pool_pre_ping') is False
|
|
|
|
|
|
def test_database_session_service_creates_read_only_engine_for_spanner():
|
|
captured_binds = []
|
|
fake_engine = mock.Mock()
|
|
fake_engine.dialect.name = 'spanner+spanner'
|
|
fake_engine.sync_engine = mock.Mock()
|
|
read_only_engine = mock.Mock()
|
|
fake_engine.execution_options.return_value = read_only_engine
|
|
|
|
def fake_async_sessionmaker(*, bind, expire_on_commit, **kwargs):
|
|
del expire_on_commit
|
|
del kwargs
|
|
captured_binds.append(bind)
|
|
return mock.Mock()
|
|
|
|
with (
|
|
mock.patch.object(
|
|
database_session_service,
|
|
'create_async_engine',
|
|
return_value=fake_engine,
|
|
),
|
|
mock.patch.object(
|
|
database_session_service,
|
|
'async_sessionmaker',
|
|
side_effect=fake_async_sessionmaker,
|
|
),
|
|
):
|
|
database_session_service.DatabaseSessionService(
|
|
'spanner+spanner:///projects/test/instances/test/databases/test'
|
|
)
|
|
|
|
assert captured_binds == [fake_engine, read_only_engine]
|
|
fake_engine.execution_options.assert_called_once_with(read_only=True)
|
|
|
|
|
|
def test_database_session_service_creates_read_only_engine_for_other_dialects():
|
|
captured_binds = []
|
|
fake_engine = mock.Mock()
|
|
fake_engine.dialect.name = 'postgresql'
|
|
fake_engine.sync_engine = mock.Mock()
|
|
read_only_engine = mock.Mock()
|
|
fake_engine.execution_options.return_value = read_only_engine
|
|
|
|
def fake_async_sessionmaker(*, bind, expire_on_commit, **kwargs):
|
|
del expire_on_commit
|
|
del kwargs
|
|
captured_binds.append(bind)
|
|
return mock.Mock()
|
|
|
|
with (
|
|
mock.patch.object(
|
|
database_session_service,
|
|
'create_async_engine',
|
|
return_value=fake_engine,
|
|
),
|
|
mock.patch.object(
|
|
database_session_service,
|
|
'async_sessionmaker',
|
|
side_effect=fake_async_sessionmaker,
|
|
),
|
|
):
|
|
database_session_service.DatabaseSessionService(
|
|
'postgresql+psycopg2://user:pass@localhost:5432/db'
|
|
)
|
|
|
|
assert captured_binds == [fake_engine, read_only_engine]
|
|
fake_engine.execution_options.assert_called_once_with(read_only=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sqlite_session_service_accepts_sqlite_urls(
|
|
tmp_path, monkeypatch
|
|
):
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
service = SqliteSessionService('sqlite+aiosqlite:///./sessions.db')
|
|
await service.create_session(app_name='app', user_id='user')
|
|
assert (tmp_path / 'sessions.db').exists()
|
|
|
|
service = SqliteSessionService('sqlite:///./sessions2.db')
|
|
await service.create_session(app_name='app', user_id='user')
|
|
assert (tmp_path / 'sessions2.db').exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sqlite_session_service_preserves_uri_query_parameters(
|
|
tmp_path, monkeypatch
|
|
):
|
|
monkeypatch.chdir(tmp_path)
|
|
db_path = tmp_path / 'readonly.db'
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.execute('CREATE TABLE IF NOT EXISTS t (id INTEGER)')
|
|
conn.commit()
|
|
|
|
service = SqliteSessionService(f'sqlite+aiosqlite:///{db_path}?mode=ro')
|
|
# `mode=ro` opens the DB read-only; schema creation should fail.
|
|
with pytest.raises(sqlite3.OperationalError, match=r'readonly'):
|
|
await service.create_session(app_name='app', user_id='user')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sqlite_session_service_accepts_absolute_sqlite_urls(tmp_path):
|
|
abs_db_path = tmp_path / 'absolute.db'
|
|
abs_url = 'sqlite+aiosqlite:////' + str(abs_db_path).lstrip('/')
|
|
service = SqliteSessionService(abs_url)
|
|
await service.create_session(app_name='app', user_id='user')
|
|
assert abs_db_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_empty_session(session_service):
|
|
assert not await session_service.get_session(
|
|
app_name='my_app', user_id='test_user', session_id='123'
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_database_session_service_get_session_uses_read_only_factory():
|
|
async with DatabaseSessionService('sqlite+aiosqlite:///:memory:') as service:
|
|
service.prepare_tables = mock.AsyncMock()
|
|
|
|
read_only_session = mock.AsyncMock()
|
|
read_only_session.get = mock.AsyncMock(return_value=None)
|
|
|
|
@asynccontextmanager
|
|
async def fake_read_only_session():
|
|
yield read_only_session
|
|
|
|
service.database_session_factory = mock.Mock(
|
|
side_effect=AssertionError('write session factory should not be used')
|
|
)
|
|
service._read_only_database_session_factory = mock.Mock(
|
|
return_value=fake_read_only_session()
|
|
)
|
|
|
|
session = await service.get_session(
|
|
app_name='my_app', user_id='test_user', session_id='123'
|
|
)
|
|
|
|
assert session is None
|
|
service._read_only_database_session_factory.assert_called_once_with()
|
|
service.database_session_factory.assert_not_called()
|
|
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_database_session_service_list_sessions_uses_read_only_factory():
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
service.prepare_tables = mock.AsyncMock()
|
|
|
|
read_only_session = mock.AsyncMock()
|
|
empty_result = mock.Mock()
|
|
empty_result.scalars.return_value.all.return_value = []
|
|
read_only_session.execute = mock.AsyncMock(return_value=empty_result)
|
|
read_only_session.get = mock.AsyncMock(return_value=None)
|
|
|
|
@asynccontextmanager
|
|
async def fake_read_only_session():
|
|
yield read_only_session
|
|
|
|
service.database_session_factory = mock.Mock(
|
|
side_effect=AssertionError('write session factory should not be used')
|
|
)
|
|
service._read_only_database_session_factory = mock.Mock(
|
|
return_value=fake_read_only_session()
|
|
)
|
|
|
|
response = await service.list_sessions(app_name='my_app', user_id='test_user')
|
|
|
|
assert response.sessions == []
|
|
service._read_only_database_session_factory.assert_called_once_with()
|
|
service.database_session_factory.assert_not_called()
|
|
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_get_session(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'test_user'
|
|
state = {'key': 'value'}
|
|
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id, state=state
|
|
)
|
|
assert session.app_name == app_name
|
|
assert session.user_id == user_id
|
|
assert session.id
|
|
assert session.state == state
|
|
assert (
|
|
session.last_update_time
|
|
<= datetime.now().astimezone(timezone.utc).timestamp()
|
|
)
|
|
|
|
got_session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
assert got_session == session
|
|
assert (
|
|
got_session.last_update_time
|
|
<= datetime.now().astimezone(timezone.utc).timestamp()
|
|
)
|
|
|
|
session_id = session.id
|
|
await session_service.delete_session(
|
|
app_name=app_name, user_id=user_id, session_id=session_id
|
|
)
|
|
|
|
assert (
|
|
await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_and_list_sessions(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'test_user'
|
|
|
|
session_ids = ['session' + str(i) for i in range(5)]
|
|
for session_id in session_ids:
|
|
await session_service.create_session(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
state={'key': 'value' + session_id},
|
|
)
|
|
|
|
list_sessions_response = await session_service.list_sessions(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
sessions = list_sessions_response.sessions
|
|
assert len(sessions) == len(session_ids)
|
|
assert {s.id for s in sessions} == set(session_ids)
|
|
for session in sessions:
|
|
assert session.state == {'key': 'value' + session.id}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_sessions_all_users(session_service):
|
|
app_name = 'my_app'
|
|
user_id_1 = 'user1'
|
|
user_id_2 = 'user2'
|
|
|
|
await session_service.create_session(
|
|
app_name=app_name,
|
|
user_id=user_id_1,
|
|
session_id='session1a',
|
|
state={'key': 'value1a'},
|
|
)
|
|
await session_service.create_session(
|
|
app_name=app_name,
|
|
user_id=user_id_1,
|
|
session_id='session1b',
|
|
state={'key': 'value1b'},
|
|
)
|
|
await session_service.create_session(
|
|
app_name=app_name,
|
|
user_id=user_id_2,
|
|
session_id='session2a',
|
|
state={'key': 'value2a'},
|
|
)
|
|
|
|
# List sessions for user1 - should contain merged state
|
|
list_sessions_response_1 = await session_service.list_sessions(
|
|
app_name=app_name, user_id=user_id_1
|
|
)
|
|
sessions_1 = list_sessions_response_1.sessions
|
|
assert len(sessions_1) == 2
|
|
sessions_1_map = {s.id: s for s in sessions_1}
|
|
assert sessions_1_map['session1a'].state == {'key': 'value1a'}
|
|
assert sessions_1_map['session1b'].state == {'key': 'value1b'}
|
|
|
|
# List sessions for user2 - should contain merged state
|
|
list_sessions_response_2 = await session_service.list_sessions(
|
|
app_name=app_name, user_id=user_id_2
|
|
)
|
|
sessions_2 = list_sessions_response_2.sessions
|
|
assert len(sessions_2) == 1
|
|
assert sessions_2[0].id == 'session2a'
|
|
assert sessions_2[0].state == {'key': 'value2a'}
|
|
|
|
# List sessions for all users - should contain merged state
|
|
list_sessions_response_all = await session_service.list_sessions(
|
|
app_name=app_name, user_id=None
|
|
)
|
|
sessions_all = list_sessions_response_all.sessions
|
|
assert len(sessions_all) == 3
|
|
sessions_all_map = {s.id: s for s in sessions_all}
|
|
assert sessions_all_map['session1a'].state == {'key': 'value1a'}
|
|
assert sessions_all_map['session1b'].state == {'key': 'value1b'}
|
|
assert sessions_all_map['session2a'].state == {'key': 'value2a'}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_state_is_shared_by_all_users_of_app(session_service):
|
|
app_name = 'my_app'
|
|
# User 1 creates a session, establishing app:k1
|
|
session1 = await session_service.create_session(
|
|
app_name=app_name, user_id='u1', session_id='s1', state={'app:k1': 'v1'}
|
|
)
|
|
# User 1 appends an event to session1, establishing app:k2
|
|
event = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
actions=EventActions(state_delta={'app:k2': 'v2'}),
|
|
)
|
|
await session_service.append_event(session=session1, event=event)
|
|
|
|
# User 2 creates a new session session2, it should see app:k1 and app:k2
|
|
session2 = await session_service.create_session(
|
|
app_name=app_name, user_id='u2', session_id='s2'
|
|
)
|
|
assert session2.state == {'app:k1': 'v1', 'app:k2': 'v2'}
|
|
|
|
# If we get session session1 again, it should also see both
|
|
session1_got = await session_service.get_session(
|
|
app_name=app_name, user_id='u1', session_id='s1'
|
|
)
|
|
assert session1_got.state.get('app:k1') == 'v1'
|
|
assert session1_got.state.get('app:k2') == 'v2'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_state_is_shared_only_by_user_sessions(session_service):
|
|
app_name = 'my_app'
|
|
# User 1 creates a session, establishing user:k1 for user 1
|
|
session1 = await session_service.create_session(
|
|
app_name=app_name, user_id='u1', session_id='s1', state={'user:k1': 'v1'}
|
|
)
|
|
# User 1 appends an event to session1, establishing user:k2 for user 1
|
|
event = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
actions=EventActions(state_delta={'user:k2': 'v2'}),
|
|
)
|
|
await session_service.append_event(session=session1, event=event)
|
|
|
|
# Another session for User 1 should see user:k1 and user:k2
|
|
session1b = await session_service.create_session(
|
|
app_name=app_name, user_id='u1', session_id='s1b'
|
|
)
|
|
assert session1b.state == {'user:k1': 'v1', 'user:k2': 'v2'}
|
|
|
|
# A session for User 2 should NOT see user:k1 or user:k2
|
|
session2 = await session_service.create_session(
|
|
app_name=app_name, user_id='u2', session_id='s2'
|
|
)
|
|
assert session2.state == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_session_state_is_not_shared(session_service):
|
|
app_name = 'my_app'
|
|
# User 1 creates a session session1, establishing sk1 only for session1
|
|
session1 = await session_service.create_session(
|
|
app_name=app_name, user_id='u1', session_id='s1', state={'sk1': 'v1'}
|
|
)
|
|
# User 1 appends an event to session1, establishing sk2 only for session1
|
|
event = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
actions=EventActions(state_delta={'sk2': 'v2'}),
|
|
)
|
|
await session_service.append_event(session=session1, event=event)
|
|
|
|
# Getting session1 should show sk1 and sk2
|
|
session1_got = await session_service.get_session(
|
|
app_name=app_name, user_id='u1', session_id='s1'
|
|
)
|
|
assert session1_got.state.get('sk1') == 'v1'
|
|
assert session1_got.state.get('sk2') == 'v2'
|
|
|
|
# Creating another session session1b for User 1 should NOT see sk1 or sk2
|
|
session1b = await session_service.create_session(
|
|
app_name=app_name, user_id='u1', session_id='s1b'
|
|
)
|
|
assert session1b.state == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_temp_state_is_not_persisted_in_state_or_events(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'u1'
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id, session_id='s1'
|
|
)
|
|
event = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
actions=EventActions(state_delta={'temp:k1': 'v1', 'sk': 'v2'}),
|
|
)
|
|
await session_service.append_event(session=session, event=event)
|
|
|
|
# Temp state IS available in the in-memory session (same invocation)
|
|
assert session.state.get('temp:k1') == 'v1'
|
|
assert session.state.get('sk') == 'v2'
|
|
|
|
# Check event as stored in session does not contain temp keys in state_delta
|
|
assert 'temp:k1' not in event.actions.state_delta
|
|
assert event.actions.state_delta.get('sk') == 'v2'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_temp_state_visible_across_sequential_events(session_service):
|
|
"""Temp state set by one event should be readable before the next event.
|
|
|
|
This simulates a SequentialAgent where agent-1 writes output_key='temp:out'
|
|
and agent-2 needs to read it from session.state within the same invocation.
|
|
"""
|
|
app_name = 'my_app'
|
|
user_id = 'u1'
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id, session_id='s_seq'
|
|
)
|
|
|
|
# Agent-1 writes temp state
|
|
event1 = Event(
|
|
invocation_id='inv1',
|
|
author='agent1',
|
|
actions=EventActions(state_delta={'temp:output': 'result_from_a1'}),
|
|
)
|
|
await session_service.append_event(session=session, event=event1)
|
|
|
|
# Agent-2 should be able to read temp state from the same session object
|
|
assert session.state.get('temp:output') == 'result_from_a1'
|
|
|
|
# But the event delta should NOT contain the temp key (not persisted)
|
|
assert 'temp:output' not in event1.actions.state_delta
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_session_respects_user_id(session_service):
|
|
app_name = 'my_app'
|
|
# u1 creates session 's1' and adds an event
|
|
session1 = await session_service.create_session(
|
|
app_name=app_name, user_id='u1', session_id='s1'
|
|
)
|
|
event = Event(invocation_id='inv1', author='user')
|
|
await session_service.append_event(session1, event)
|
|
# u2 creates a session with the same session_id 's1'
|
|
await session_service.create_session(
|
|
app_name=app_name, user_id='u2', session_id='s1'
|
|
)
|
|
# Check that getting s1 for u2 returns u2's session (with no events)
|
|
# not u1's session.
|
|
session2_got = await session_service.get_session(
|
|
app_name=app_name, user_id='u2', session_id='s1'
|
|
)
|
|
assert session2_got.user_id == 'u2'
|
|
assert len(session2_got.events) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_session_with_existing_id_raises_error(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'test_user'
|
|
session_id = 'existing_session'
|
|
|
|
# Create the first session
|
|
await session_service.create_session(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
)
|
|
|
|
# Attempt to create a session with the same ID
|
|
with pytest.raises(AlreadyExistsError):
|
|
await session_service.create_session(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_bytes(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'user'
|
|
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
|
|
test_content = types.Content(
|
|
role='user',
|
|
parts=[
|
|
types.Part.from_bytes(data=b'test_image_data', mime_type='image/png'),
|
|
],
|
|
)
|
|
test_grounding_metadata = types.GroundingMetadata(
|
|
search_entry_point=types.SearchEntryPoint(sdk_blob=b'test_sdk_blob')
|
|
)
|
|
event = Event(
|
|
invocation_id='invocation',
|
|
author='user',
|
|
content=test_content,
|
|
grounding_metadata=test_grounding_metadata,
|
|
)
|
|
await session_service.append_event(session=session, event=event)
|
|
|
|
assert session.events[0].content == test_content
|
|
|
|
session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
events = session.events
|
|
assert len(events) == 1
|
|
assert events[0].content == test_content
|
|
assert events[0].grounding_metadata == test_grounding_metadata
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_complete(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'user'
|
|
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
event = Event(
|
|
invocation_id='invocation',
|
|
author='user',
|
|
content=types.Content(role='user', parts=[types.Part(text='test_text')]),
|
|
turn_complete=True,
|
|
partial=False,
|
|
actions=EventActions(
|
|
artifact_delta={
|
|
'file': 0,
|
|
},
|
|
transfer_to_agent='agent',
|
|
escalate=True,
|
|
),
|
|
long_running_tool_ids={'tool1'},
|
|
error_code='error_code',
|
|
error_message='error_message',
|
|
interrupted=True,
|
|
grounding_metadata=types.GroundingMetadata(
|
|
web_search_queries=['query1'],
|
|
),
|
|
usage_metadata=types.GenerateContentResponseUsageMetadata(
|
|
prompt_token_count=1, candidates_token_count=1, total_token_count=2
|
|
),
|
|
citation_metadata=types.CitationMetadata(),
|
|
custom_metadata={'custom_key': 'custom_value'},
|
|
timestamp=1700000000.123,
|
|
input_transcription=types.Transcription(
|
|
text='input transcription',
|
|
finished=True,
|
|
),
|
|
output_transcription=types.Transcription(
|
|
text='output transcription',
|
|
finished=True,
|
|
),
|
|
)
|
|
await session_service.append_event(session=session, event=event)
|
|
|
|
assert (
|
|
await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
== session
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_session_last_update_time_updates_on_event(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'user'
|
|
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
original_update_time = session.last_update_time
|
|
|
|
event_timestamp = original_update_time + 10
|
|
event = Event(
|
|
invocation_id='invocation',
|
|
author='user',
|
|
timestamp=event_timestamp,
|
|
)
|
|
await session_service.append_event(session=session, event=event)
|
|
|
|
assert session.last_update_time == pytest.approx(event_timestamp, abs=1e-6)
|
|
|
|
refreshed_session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
assert refreshed_session is not None
|
|
assert refreshed_session.last_update_time == pytest.approx(
|
|
event_timestamp, abs=1e-6
|
|
)
|
|
assert refreshed_session.last_update_time > original_update_time
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_to_stale_session():
|
|
session_service = get_session_service(
|
|
service_type=SessionServiceType.DATABASE
|
|
)
|
|
|
|
async with session_service:
|
|
app_name = 'my_app'
|
|
user_id = 'user'
|
|
current_time = datetime.now().astimezone(timezone.utc).timestamp()
|
|
|
|
original_session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
event1 = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
timestamp=current_time + 1,
|
|
actions=EventActions(state_delta={'sk1': 'v1'}),
|
|
)
|
|
await session_service.append_event(original_session, event1)
|
|
|
|
updated_session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=original_session.id
|
|
)
|
|
event2 = Event(
|
|
invocation_id='inv2',
|
|
author='user',
|
|
timestamp=current_time + 2,
|
|
actions=EventActions(state_delta={'sk2': 'v2'}),
|
|
)
|
|
await session_service.append_event(updated_session, event2)
|
|
|
|
# original_session is now stale
|
|
assert original_session.last_update_time < updated_session.last_update_time
|
|
assert len(original_session.events) == 1
|
|
assert 'sk2' not in original_session.state
|
|
|
|
# Appending another event to stale original_session should be rejected.
|
|
event3 = Event(
|
|
invocation_id='inv3',
|
|
author='user',
|
|
timestamp=current_time + 3,
|
|
actions=EventActions(state_delta={'sk3': 'v3'}),
|
|
)
|
|
with pytest.raises(ValueError, match='modified in storage'):
|
|
await session_service.append_event(original_session, event3)
|
|
|
|
# If we fetch session from DB, it should only contain the committed events.
|
|
session_final = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=original_session.id
|
|
)
|
|
assert len(session_final.events) == 2
|
|
assert session_final.state.get('sk1') == 'v1'
|
|
assert session_final.state.get('sk2') == 'v2'
|
|
assert session_final.state.get('sk3') is None
|
|
assert [e.invocation_id for e in session_final.events] == [
|
|
'inv1',
|
|
'inv2',
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_raises_if_app_state_row_missing():
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
session = await service.create_session(
|
|
app_name='my_app', user_id='user', session_id='s1'
|
|
)
|
|
schema = service._get_schema_classes()
|
|
async with service.database_session_factory() as sql_session:
|
|
await sql_session.execute(
|
|
delete(schema.StorageAppState).where(
|
|
schema.StorageAppState.app_name == session.app_name
|
|
)
|
|
)
|
|
await sql_session.commit()
|
|
|
|
event = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
actions=EventActions(state_delta={'k': 'v'}),
|
|
)
|
|
with pytest.raises(ValueError, match='App state missing'):
|
|
await service.append_event(session, event)
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_raises_if_user_state_row_missing():
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
session = await service.create_session(
|
|
app_name='my_app', user_id='user', session_id='s1'
|
|
)
|
|
schema = service._get_schema_classes()
|
|
async with service.database_session_factory() as sql_session:
|
|
await sql_session.execute(
|
|
delete(schema.StorageUserState).where(
|
|
schema.StorageUserState.app_name == session.app_name,
|
|
schema.StorageUserState.user_id == session.user_id,
|
|
)
|
|
)
|
|
await sql_session.commit()
|
|
|
|
event = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
actions=EventActions(state_delta={'k': 'v'}),
|
|
)
|
|
with pytest.raises(ValueError, match='User state missing'):
|
|
await service.append_event(session, event)
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_concurrent_stale_sessions_reject_stale_writer():
|
|
session_service = get_session_service(
|
|
service_type=SessionServiceType.DATABASE
|
|
)
|
|
|
|
async with session_service:
|
|
app_name = 'my_app'
|
|
user_id = 'user'
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
|
|
iteration_count = 8
|
|
for i in range(iteration_count):
|
|
latest_session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
stale_session_1 = latest_session.model_copy(deep=True)
|
|
stale_session_2 = latest_session.model_copy(deep=True)
|
|
base_timestamp = latest_session.last_update_time + 10.0
|
|
event_1 = Event(
|
|
invocation_id=f'inv{i}-1',
|
|
author='user',
|
|
timestamp=base_timestamp + 1.0,
|
|
actions=EventActions(state_delta={f'sk{i}-1': f'v{i}-1'}),
|
|
)
|
|
event_2 = Event(
|
|
invocation_id=f'inv{i}-2',
|
|
author='user',
|
|
timestamp=base_timestamp + 2.0,
|
|
actions=EventActions(state_delta={f'sk{i}-2': f'v{i}-2'}),
|
|
)
|
|
|
|
results = await asyncio.gather(
|
|
session_service.append_event(stale_session_1, event_1),
|
|
session_service.append_event(stale_session_2, event_2),
|
|
return_exceptions=True,
|
|
)
|
|
errors = [result for result in results if isinstance(result, Exception)]
|
|
successes = [
|
|
result for result in results if not isinstance(result, Exception)
|
|
]
|
|
assert len(successes) == 1
|
|
assert len(errors) == 1
|
|
assert isinstance(errors[0], ValueError)
|
|
assert 'modified in storage' in str(errors[0])
|
|
|
|
session_final = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
|
|
for i in range(iteration_count):
|
|
event_values = {
|
|
session_final.state.get(f'sk{i}-1'),
|
|
session_final.state.get(f'sk{i}-2'),
|
|
}
|
|
assert event_values & {f'v{i}-1', f'v{i}-2'}
|
|
assert None in event_values
|
|
assert len(session_final.events) == iteration_count
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_allows_timestamp_drift_for_current_session():
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
session = await service.create_session(
|
|
app_name='my_app', user_id='user', session_id='s1'
|
|
)
|
|
event1 = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
timestamp=session.last_update_time + 10,
|
|
)
|
|
await service.append_event(session, event1)
|
|
|
|
# Simulate a float round-trip mismatch without changing the persisted
|
|
# session revision.
|
|
session.last_update_time -= 0.0001
|
|
|
|
event2 = Event(
|
|
invocation_id='inv2',
|
|
author='user',
|
|
timestamp=event1.timestamp + 10,
|
|
)
|
|
await service.append_event(session, event2)
|
|
|
|
refreshed_session = await service.get_session(
|
|
app_name='my_app', user_id='user', session_id=session.id
|
|
)
|
|
assert [event.invocation_id for event in refreshed_session.events] == [
|
|
'inv1',
|
|
'inv2',
|
|
]
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_allows_markerless_current_session():
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
session = await service.create_session(
|
|
app_name='my_app', user_id='user', session_id='s1'
|
|
)
|
|
event1 = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
timestamp=session.last_update_time + 10,
|
|
)
|
|
await service.append_event(session, event1)
|
|
|
|
session._storage_update_marker = None
|
|
session.last_update_time -= 0.0001
|
|
|
|
event2 = Event(
|
|
invocation_id='inv2',
|
|
author='user',
|
|
timestamp=event1.timestamp + 10,
|
|
)
|
|
await service.append_event(session, event2)
|
|
|
|
refreshed_session = await service.get_session(
|
|
app_name='my_app', user_id='user', session_id=session.id
|
|
)
|
|
assert [event.invocation_id for event in refreshed_session.events] == [
|
|
'inv1',
|
|
'inv2',
|
|
]
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_when_session_is_same_ref_as_storage_session():
|
|
"""Tests that appending an event to a session only appends it once if the user-passed session and the underlying storage session are the same object."""
|
|
service = InMemorySessionService()
|
|
app_name = 'my_app'
|
|
user_id = 'test_user'
|
|
|
|
# Create a session
|
|
session = await service.create_session(app_name=app_name, user_id=user_id)
|
|
|
|
# Get the actual storage event object from the underlying storage
|
|
storage_session = service.sessions[app_name][user_id][session.id]
|
|
|
|
# Append the event to the storage session directly
|
|
event = Event(invocation_id='inv1', author='user')
|
|
await service.append_event(session=storage_session, event=event)
|
|
|
|
# Verify that the storage session has only one event
|
|
final_session = await service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
assert len(final_session.events) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_session_with_config(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'user'
|
|
|
|
num_test_events = 5
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
for i in range(1, num_test_events + 1):
|
|
event = Event(author='user', timestamp=i)
|
|
await session_service.append_event(session, event)
|
|
|
|
# No config, expect all events to be returned.
|
|
session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
events = session.events
|
|
assert len(events) == num_test_events
|
|
|
|
# Explicitly requesting zero recent events should return no event history.
|
|
config = GetSessionConfig(num_recent_events=0)
|
|
session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id, config=config
|
|
)
|
|
assert not session.events
|
|
|
|
# Only expect the most recent 3 events.
|
|
num_recent_events = 3
|
|
config = GetSessionConfig(num_recent_events=num_recent_events)
|
|
session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id, config=config
|
|
)
|
|
events = session.events
|
|
assert len(events) == num_recent_events
|
|
assert events[0].timestamp == num_test_events - num_recent_events + 1
|
|
|
|
# Only expect events after timestamp 4.0 (inclusive), i.e., 2 events.
|
|
after_timestamp = 4.0
|
|
config = GetSessionConfig(after_timestamp=after_timestamp)
|
|
session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id, config=config
|
|
)
|
|
events = session.events
|
|
assert len(events) == num_test_events - after_timestamp + 1
|
|
assert events[0].timestamp == after_timestamp
|
|
|
|
# Expect no events if none are > after_timestamp.
|
|
way_after_timestamp = num_test_events * 10
|
|
config = GetSessionConfig(after_timestamp=way_after_timestamp)
|
|
session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id, config=config
|
|
)
|
|
assert not session.events
|
|
|
|
# Both filters applied, i.e., of 3 most recent events, only 2 are after
|
|
# timestamp 4.0, so expect 2 events.
|
|
config = GetSessionConfig(
|
|
after_timestamp=after_timestamp, num_recent_events=num_recent_events
|
|
)
|
|
session = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id, config=config
|
|
)
|
|
events = session.events
|
|
assert len(events) == num_test_events - after_timestamp + 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_partial_events_are_not_persisted(session_service):
|
|
app_name = 'my_app'
|
|
user_id = 'user'
|
|
session = await session_service.create_session(
|
|
app_name=app_name, user_id=user_id
|
|
)
|
|
event = Event(author='user', partial=True)
|
|
await session_service.append_event(session, event)
|
|
|
|
# Check in-memory session
|
|
assert len(session.events) == 0
|
|
# Check persisted session
|
|
session_got = await session_service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
assert len(session_got.events) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rollback tests – verify _rollback_on_exception_session explicitly rolls back
|
|
# on errors
|
|
# ---------------------------------------------------------------------------
|
|
class _RollbackSpySession:
|
|
"""Wraps an AsyncSession to spy on rollback() and optionally fail commit()."""
|
|
|
|
def __init__(self, real_session, *, fail_commit=False):
|
|
self._real = real_session
|
|
self._fail_commit = fail_commit
|
|
self.rollback_called = False
|
|
|
|
async def __aenter__(self):
|
|
self._real = await self._real.__aenter__()
|
|
return self
|
|
|
|
async def __aexit__(self, *args):
|
|
return await self._real.__aexit__(*args)
|
|
|
|
async def commit(self):
|
|
if self._fail_commit:
|
|
raise RuntimeError('simulated commit failure')
|
|
return await self._real.commit()
|
|
|
|
async def rollback(self):
|
|
self.rollback_called = True
|
|
return await self._real.rollback()
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._real, name)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_session_calls_rollback_on_commit_failure():
|
|
"""Verifies that a commit failure during create_session triggers an explicit
|
|
rollback() call via _rollback_on_exception_session, not just a close()."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
# Ensure tables are initialized.
|
|
await service.create_session(
|
|
app_name='app', user_id='user', session_id='good'
|
|
)
|
|
|
|
original_factory = service.database_session_factory
|
|
spy_sessions = []
|
|
|
|
def _spy_factory():
|
|
spy = _RollbackSpySession(original_factory(), fail_commit=True)
|
|
spy_sessions.append(spy)
|
|
return spy
|
|
|
|
service.database_session_factory = _spy_factory
|
|
|
|
with pytest.raises(RuntimeError, match='simulated commit failure'):
|
|
await service.create_session(
|
|
app_name='app', user_id='user', session_id='should_fail'
|
|
)
|
|
|
|
# The key assertion: rollback() must have been called explicitly.
|
|
assert len(spy_sessions) == 1
|
|
assert spy_sessions[0].rollback_called, (
|
|
'rollback() was not called – _rollback_on_exception_session is not'
|
|
' protecting this path'
|
|
)
|
|
|
|
# Restore and verify the failed session was not persisted.
|
|
service.database_session_factory = original_factory
|
|
assert (
|
|
await service.get_session(
|
|
app_name='app', user_id='user', session_id='should_fail'
|
|
)
|
|
is None
|
|
)
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_calls_rollback_on_commit_failure():
|
|
"""Verifies that a commit failure during append_event triggers an explicit
|
|
rollback() call via _rollback_on_exception_session."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
|
|
# Successfully append one event first.
|
|
event1 = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
actions=EventActions(state_delta={'key1': 'value1'}),
|
|
)
|
|
await service.append_event(session, event1)
|
|
|
|
original_factory = service.database_session_factory
|
|
spy_sessions = []
|
|
|
|
def _spy_factory():
|
|
spy = _RollbackSpySession(original_factory(), fail_commit=True)
|
|
spy_sessions.append(spy)
|
|
return spy
|
|
|
|
service.database_session_factory = _spy_factory
|
|
|
|
event2 = Event(
|
|
invocation_id='inv2',
|
|
author='user',
|
|
actions=EventActions(state_delta={'key2': 'value2'}),
|
|
)
|
|
with pytest.raises(RuntimeError, match='simulated commit failure'):
|
|
await service.append_event(session, event2)
|
|
|
|
assert len(spy_sessions) == 1
|
|
assert spy_sessions[0].rollback_called, (
|
|
'rollback() was not called – _rollback_on_exception_session is not'
|
|
' protecting this path'
|
|
)
|
|
|
|
# Restore and verify only the first event was persisted.
|
|
service.database_session_factory = original_factory
|
|
got = await service.get_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
assert len(got.events) == 1
|
|
assert got.events[0].invocation_id == 'inv1'
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
class _CommitOrderSpySession:
|
|
"""SQLAlchemy session spy that marks when commit() has completed."""
|
|
|
|
def __init__(self, real_session, on_committed):
|
|
self._real = real_session
|
|
self._on_committed = on_committed
|
|
|
|
async def __aenter__(self):
|
|
self._real = await self._real.__aenter__()
|
|
return self
|
|
|
|
async def __aexit__(self, *args):
|
|
return await self._real.__aexit__(*args)
|
|
|
|
async def commit(self):
|
|
result = await self._real.commit()
|
|
self._on_committed()
|
|
return result
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._real, name)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_event_reads_storage_revision_before_commit():
|
|
"""append_event captures session revision before commit completes."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
await service.prepare_tables()
|
|
schema = service._get_schema_classes()
|
|
original_get_update_timestamp = schema.StorageSession.get_update_timestamp
|
|
original_get_update_marker = schema.StorageSession.get_update_marker
|
|
revision_read_state = {'committed': False, 'post_commit_reads': 0}
|
|
|
|
def _track_revision_read(original):
|
|
def wrapper(self, *args, **kwargs):
|
|
if revision_read_state['committed']:
|
|
revision_read_state['post_commit_reads'] += 1
|
|
return original(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
schema.StorageSession.get_update_timestamp = _track_revision_read(
|
|
original_get_update_timestamp
|
|
)
|
|
schema.StorageSession.get_update_marker = _track_revision_read(
|
|
original_get_update_marker
|
|
)
|
|
|
|
try:
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
event_timestamp = session.last_update_time + 10
|
|
event = Event(
|
|
invocation_id='inv1',
|
|
author='user',
|
|
timestamp=event_timestamp,
|
|
)
|
|
|
|
original_factory = service.database_session_factory
|
|
|
|
def _spy_factory():
|
|
return _CommitOrderSpySession(
|
|
original_factory(),
|
|
on_committed=lambda: revision_read_state.update({'committed': True}),
|
|
)
|
|
|
|
service.database_session_factory = _spy_factory
|
|
|
|
await service.append_event(session, event)
|
|
|
|
assert revision_read_state['post_commit_reads'] == 0
|
|
assert session.last_update_time == pytest.approx(event_timestamp, abs=1e-6)
|
|
assert session._storage_update_marker is not None
|
|
finally:
|
|
schema.StorageSession.get_update_timestamp = original_get_update_timestamp
|
|
schema.StorageSession.get_update_marker = original_get_update_marker
|
|
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_session_reads_storage_revision_before_commit():
|
|
"""create_session captures session revision before commit completes."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
await service.prepare_tables()
|
|
schema = service._get_schema_classes()
|
|
original_get_update_timestamp = schema.StorageSession.get_update_timestamp
|
|
original_get_update_marker = schema.StorageSession.get_update_marker
|
|
revision_read_state = {'committed': False, 'post_commit_reads': 0}
|
|
|
|
def _track_revision_read(original):
|
|
def wrapper(self, *args, **kwargs):
|
|
if revision_read_state['committed']:
|
|
revision_read_state['post_commit_reads'] += 1
|
|
return original(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
schema.StorageSession.get_update_timestamp = _track_revision_read(
|
|
original_get_update_timestamp
|
|
)
|
|
schema.StorageSession.get_update_marker = _track_revision_read(
|
|
original_get_update_marker
|
|
)
|
|
|
|
try:
|
|
original_factory = service.database_session_factory
|
|
|
|
def _spy_factory():
|
|
return _CommitOrderSpySession(
|
|
original_factory(),
|
|
on_committed=lambda: revision_read_state.update({'committed': True}),
|
|
)
|
|
|
|
service.database_session_factory = _spy_factory
|
|
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
|
|
assert revision_read_state['post_commit_reads'] == 0
|
|
assert session.last_update_time is not None
|
|
assert session._storage_update_marker is not None
|
|
finally:
|
|
schema.StorageSession.get_update_timestamp = original_get_update_timestamp
|
|
schema.StorageSession.get_update_marker = original_get_update_marker
|
|
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_session_calls_rollback_on_commit_failure():
|
|
"""Verifies that a commit failure during delete_session triggers an explicit
|
|
rollback() call via _rollback_on_exception_session."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
|
|
original_factory = service.database_session_factory
|
|
spy_sessions = []
|
|
|
|
def _spy_factory():
|
|
spy = _RollbackSpySession(original_factory(), fail_commit=True)
|
|
spy_sessions.append(spy)
|
|
return spy
|
|
|
|
service.database_session_factory = _spy_factory
|
|
|
|
with pytest.raises(RuntimeError, match='simulated commit failure'):
|
|
await service.delete_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
|
|
assert len(spy_sessions) == 1
|
|
assert spy_sessions[0].rollback_called, (
|
|
'rollback() was not called – _rollback_on_exception_session is not'
|
|
' protecting this path'
|
|
)
|
|
|
|
# Restore and verify the session still exists (delete was rolled back).
|
|
service.database_session_factory = original_factory
|
|
got = await service.get_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
assert got is not None
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_recovers_after_multiple_failures():
|
|
"""After several consecutive commit failures, every single one must trigger
|
|
a rollback() call and the service must remain functional afterward."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
await service.create_session(
|
|
app_name='app', user_id='user', session_id='seed'
|
|
)
|
|
|
|
original_factory = service.database_session_factory
|
|
spy_sessions = []
|
|
|
|
def _spy_factory():
|
|
spy = _RollbackSpySession(original_factory(), fail_commit=True)
|
|
spy_sessions.append(spy)
|
|
return spy
|
|
|
|
service.database_session_factory = _spy_factory
|
|
|
|
num_failures = 5
|
|
for i in range(num_failures):
|
|
with pytest.raises(RuntimeError, match='simulated commit failure'):
|
|
await service.create_session(
|
|
app_name='app', user_id='user', session_id=f'fail_{i}'
|
|
)
|
|
|
|
# Every failure must have triggered a rollback.
|
|
assert len(spy_sessions) == num_failures
|
|
for i, spy in enumerate(spy_sessions):
|
|
assert spy.rollback_called, f'rollback() was not called on failure #{i}'
|
|
|
|
# Restore and verify the service is still healthy.
|
|
service.database_session_factory = original_factory
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='recovered'
|
|
)
|
|
assert session.id == 'recovered'
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_prepare_tables_no_race_condition():
|
|
"""Verifies that concurrent calls to prepare_tables wait for table creation.
|
|
Reproduces the race condition where concurrent requests
|
|
arrive at startup, prepare_tables must not return before tables exist.
|
|
Previously, the early-return guard checked _db_schema_version (set during
|
|
schema detection) instead of _tables_created, so a second request could
|
|
slip through after schema detection but before table creation finished.
|
|
"""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
# Tables haven't been created yet.
|
|
assert not service._tables_created
|
|
assert service._db_schema_version is None
|
|
|
|
# Launch several concurrent create_session calls, each with a unique
|
|
# app_name to avoid IntegrityError on the shared app_states row.
|
|
# Each will call prepare_tables internally. If the race condition
|
|
# exists, some of these will fail because the "sessions" table doesn't
|
|
# exist yet.
|
|
num_concurrent = 5
|
|
results = await asyncio.gather(
|
|
*[
|
|
service.create_session(
|
|
app_name=f'app_{i}', user_id='user', session_id=f'sess_{i}'
|
|
)
|
|
for i in range(num_concurrent)
|
|
],
|
|
return_exceptions=True,
|
|
)
|
|
|
|
# Every call must succeed – no exceptions allowed.
|
|
for i, result in enumerate(results):
|
|
assert not isinstance(result, BaseException), (
|
|
f'Concurrent create_session #{i} raised {result!r}; tables were'
|
|
' likely not ready due to the prepare_tables race condition.'
|
|
)
|
|
|
|
# All sessions should be retrievable.
|
|
for i in range(num_concurrent):
|
|
session = await service.get_session(
|
|
app_name=f'app_{i}', user_id='user', session_id=f'sess_{i}'
|
|
)
|
|
assert session is not None, f'Session sess_{i} not found after creation.'
|
|
|
|
assert service._tables_created
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_tables_serializes_schema_detection_and_creation():
|
|
"""Verifies schema detection and table creation happen atomically under one
|
|
lock, so concurrent callers cannot observe a partially-initialized state.
|
|
After prepare_tables completes, both _db_schema_version and _tables_created
|
|
must be set.
|
|
"""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
assert not service._tables_created
|
|
assert service._db_schema_version is None
|
|
|
|
await service.prepare_tables()
|
|
|
|
# Both must be set after a single prepare_tables call.
|
|
assert service._tables_created
|
|
assert service._db_schema_version is not None
|
|
|
|
# Verify tables actually exist by performing a real operation.
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
assert session is not None
|
|
assert session.id == 's1'
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_or_create_state_returns_existing_row():
|
|
"""_get_or_create_state returns an existing row without inserting."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
await service.prepare_tables()
|
|
schema = service._get_schema_classes()
|
|
|
|
# Pre-create the app_state row.
|
|
async with service.database_session_factory() as sql_session:
|
|
sql_session.add(schema.StorageAppState(app_name='app1', state={'k': 'v'}))
|
|
await sql_session.commit()
|
|
|
|
# _get_or_create_state should find and return it.
|
|
async with service.database_session_factory() as sql_session:
|
|
row = await database_session_service._get_or_create_state(
|
|
sql_session=sql_session,
|
|
state_model=schema.StorageAppState,
|
|
primary_key='app1',
|
|
defaults={'app_name': 'app1', 'state': {}},
|
|
)
|
|
assert row.app_name == 'app1'
|
|
assert row.state == {'k': 'v'}
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_or_create_state_creates_new_row():
|
|
"""_get_or_create_state creates a row when none exists."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
await service.prepare_tables()
|
|
schema = service._get_schema_classes()
|
|
|
|
async with service.database_session_factory() as sql_session:
|
|
row = await database_session_service._get_or_create_state(
|
|
sql_session=sql_session,
|
|
state_model=schema.StorageAppState,
|
|
primary_key='new_app',
|
|
defaults={'app_name': 'new_app', 'state': {}},
|
|
)
|
|
await sql_session.commit()
|
|
assert row.app_name == 'new_app'
|
|
assert row.state == {}
|
|
|
|
# Verify the row was actually persisted.
|
|
async with service.database_session_factory() as sql_session:
|
|
persisted = await sql_session.get(schema.StorageAppState, 'new_app')
|
|
assert persisted is not None
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_or_create_state_handles_race_condition():
|
|
"""_get_or_create_state recovers when a concurrent INSERT wins the race.
|
|
|
|
Simulates the race:
|
|
the initial SELECT returns None (another caller hasn't committed yet), but
|
|
by the time we INSERT, the other caller has committed — so the INSERT fails
|
|
with IntegrityError and we fall back to re-fetching.
|
|
"""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
await service.prepare_tables()
|
|
schema = service._get_schema_classes()
|
|
|
|
# Pre-create the row to guarantee the INSERT will fail.
|
|
async with service.database_session_factory() as sql_session:
|
|
sql_session.add(schema.StorageAppState(app_name='race_app', state={}))
|
|
await sql_session.commit()
|
|
|
|
# Patch session.get to return None on the first call (simulating the
|
|
# race window), then fall through to the real implementation.
|
|
async with service.database_session_factory() as sql_session:
|
|
original_get = sql_session.get
|
|
call_count = 0
|
|
|
|
async def patched_get(*args, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 1:
|
|
return None # Simulate: row not yet visible
|
|
return await original_get(*args, **kwargs)
|
|
|
|
sql_session.get = patched_get
|
|
|
|
row = await database_session_service._get_or_create_state(
|
|
sql_session=sql_session,
|
|
state_model=schema.StorageAppState,
|
|
primary_key='race_app',
|
|
defaults={'app_name': 'race_app', 'state': {}},
|
|
)
|
|
assert row.app_name == 'race_app'
|
|
# The function should have called get twice: once before the INSERT
|
|
# (patched to return None) and once after the IntegrityError.
|
|
assert call_count == 2
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_session_sequential_same_app_name():
|
|
"""Sequential create_session calls for the same app_name work correctly.
|
|
|
|
The second call reuses the existing app_states row.
|
|
"""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
s1 = await service.create_session(
|
|
app_name='shared', user_id='u1', session_id='s1'
|
|
)
|
|
s2 = await service.create_session(
|
|
app_name='shared', user_id='u2', session_id='s2'
|
|
)
|
|
assert s1.app_name == 'shared'
|
|
assert s2.app_name == 'shared'
|
|
|
|
got1 = await service.get_session(
|
|
app_name='shared', user_id='u1', session_id='s1'
|
|
)
|
|
got2 = await service.get_session(
|
|
app_name='shared', user_id='u2', session_id='s2'
|
|
)
|
|
assert got1 is not None
|
|
assert got2 is not None
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_tables_idempotent_after_creation():
|
|
"""Calling prepare_tables multiple times is safe and idempotent.
|
|
After tables are created, subsequent calls should return immediately via
|
|
the fast path without errors.
|
|
"""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
try:
|
|
await service.prepare_tables()
|
|
assert service._tables_created
|
|
|
|
# Call again — should be a no-op via the fast path.
|
|
await service.prepare_tables()
|
|
assert service._tables_created
|
|
|
|
# Service should still work.
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
assert session.id == 's1'
|
|
finally:
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_public_prepare_tables_eager_initialization():
|
|
"""Calling the public prepare_tables() eagerly initializes tables so that
|
|
the first real database operation does not pay the setup cost.
|
|
"""
|
|
async with DatabaseSessionService('sqlite+aiosqlite:///:memory:') as service:
|
|
# Before calling prepare_tables, tables are not created.
|
|
assert not service._tables_created
|
|
assert service._db_schema_version is None
|
|
|
|
# Eagerly prepare tables via the public API.
|
|
await service.prepare_tables()
|
|
|
|
# Tables should now be ready.
|
|
assert service._tables_created
|
|
assert service._db_schema_version is not None
|
|
|
|
# Subsequent operations should work without any additional setup cost.
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
assert session.id == 's1'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
'state_delta, expect_app_lock, expect_user_lock',
|
|
[
|
|
pytest.param(
|
|
None,
|
|
False,
|
|
False,
|
|
id='no_state_delta',
|
|
),
|
|
pytest.param(
|
|
{'session_key': 'v'},
|
|
False,
|
|
False,
|
|
id='session_only_delta',
|
|
),
|
|
pytest.param(
|
|
{'app:key': 'v'},
|
|
True,
|
|
False,
|
|
id='app_delta_only',
|
|
),
|
|
pytest.param(
|
|
{'user:key': 'v'},
|
|
False,
|
|
True,
|
|
id='user_delta_only',
|
|
),
|
|
pytest.param(
|
|
{'app:a': '1', 'user:b': '2', 'sk': '3'},
|
|
True,
|
|
True,
|
|
id='all_scopes',
|
|
),
|
|
],
|
|
)
|
|
async def test_append_event_locks_only_scopes_with_deltas(
|
|
state_delta, expect_app_lock, expect_user_lock
|
|
):
|
|
"""FOR UPDATE should only be requested for state scopes that have deltas."""
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
|
|
lock_requests = []
|
|
original_fn = database_session_service._select_required_state
|
|
|
|
async def tracking_fn(**kwargs):
|
|
lock_requests.append({
|
|
'model': kwargs['state_model'].__tablename__,
|
|
'use_row_level_locking': kwargs['use_row_level_locking'],
|
|
})
|
|
return await original_fn(**kwargs)
|
|
|
|
try:
|
|
session = await service.create_session(
|
|
app_name='app', user_id='user', session_id='s1'
|
|
)
|
|
|
|
database_session_service._select_required_state = tracking_fn
|
|
lock_requests.clear()
|
|
|
|
event_kwargs = {'invocation_id': 'inv', 'author': 'user'}
|
|
if state_delta is not None:
|
|
event_kwargs['actions'] = EventActions(state_delta=state_delta)
|
|
event = Event(**event_kwargs)
|
|
await service.append_event(session, event)
|
|
|
|
app_req = next(
|
|
(r for r in lock_requests if r['model'] == 'app_states'), None
|
|
)
|
|
user_req = next(
|
|
(r for r in lock_requests if r['model'] == 'user_states'), None
|
|
)
|
|
|
|
# SQLite doesn't support row-level locking so use_row_level_locking is
|
|
# always False. The important check is that locking is not requested
|
|
# when there is no delta (it must never be True without a delta).
|
|
if not expect_app_lock:
|
|
assert (
|
|
app_req is None or not app_req['use_row_level_locking']
|
|
), 'app_states should not be locked without an app: delta'
|
|
if not expect_user_lock:
|
|
assert (
|
|
user_req is None or not user_req['use_row_level_locking']
|
|
), 'user_states should not be locked without a user: delta'
|
|
finally:
|
|
database_session_service._select_required_state = original_fn
|
|
await service.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_state_returns_empty_dict_when_no_state_exists(
|
|
session_service,
|
|
):
|
|
"""Verifies get_user_state returns empty dict when no state exists."""
|
|
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
|
|
assert not state
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_state_returns_state_written_via_append_event(
|
|
session_service,
|
|
):
|
|
"""Verifies get_user_state returns state written via append_event."""
|
|
session = await session_service.create_session(
|
|
app_name='my_app', user_id='u1'
|
|
)
|
|
await session_service.append_event(
|
|
session,
|
|
Event(
|
|
author='system',
|
|
actions=EventActions(
|
|
state_delta={'user:profile': {'name': 'Alice'}, 'session_key': 1}
|
|
),
|
|
),
|
|
)
|
|
|
|
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
|
|
|
|
assert state == {'profile': {'name': 'Alice'}}
|
|
assert 'session_key' not in state
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_state_is_not_visible_across_users(session_service):
|
|
"""Verifies user state is isolated between users."""
|
|
session = await session_service.create_session(
|
|
app_name='my_app', user_id='u1'
|
|
)
|
|
await session_service.append_event(
|
|
session,
|
|
Event(
|
|
author='system',
|
|
actions=EventActions(state_delta={'user:secret': 'only-for-u1'}),
|
|
),
|
|
)
|
|
|
|
other_state = await session_service.get_user_state(
|
|
app_name='my_app', user_id='u2'
|
|
)
|
|
assert not other_state
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_state_is_not_visible_across_apps(session_service):
|
|
"""Verifies user state is isolated between apps."""
|
|
session = await session_service.create_session(
|
|
app_name='my_app', user_id='u1'
|
|
)
|
|
await session_service.append_event(
|
|
session,
|
|
Event(
|
|
author='system',
|
|
actions=EventActions(state_delta={'user:data': 'only-app-a'}),
|
|
),
|
|
)
|
|
|
|
other_state = await session_service.get_user_state(
|
|
app_name='other_app', user_id='u1'
|
|
)
|
|
assert not other_state
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_state_available_before_session_is_created(
|
|
session_service,
|
|
):
|
|
"""Verifies user state can be retrieved before a session is created."""
|
|
first_session = await session_service.create_session(
|
|
app_name='my_app', user_id='u1'
|
|
)
|
|
await session_service.append_event(
|
|
first_session,
|
|
Event(
|
|
author='system',
|
|
actions=EventActions(state_delta={'user:ctx': {'v': 1}}),
|
|
),
|
|
)
|
|
|
|
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
|
|
assert state == {'ctx': {'v': 1}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_state_reflects_latest_write(session_service):
|
|
"""Verifies get_user_state returns the latest state."""
|
|
session = await session_service.create_session(
|
|
app_name='my_app', user_id='u1'
|
|
)
|
|
await session_service.append_event(
|
|
session,
|
|
Event(
|
|
author='system',
|
|
actions=EventActions(state_delta={'user:counter': 1}),
|
|
),
|
|
)
|
|
await session_service.append_event(
|
|
session,
|
|
Event(
|
|
author='system',
|
|
actions=EventActions(state_delta={'user:counter': 2}),
|
|
),
|
|
)
|
|
|
|
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
|
|
assert state['counter'] == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vertex_ai_session_service_raises_not_implemented_for_get_user_state():
|
|
"""Verifies VertexAiSessionService raises NotImplementedError."""
|
|
service = VertexAiSessionService(project='proj', location='us-central1')
|
|
with pytest.raises(NotImplementedError):
|
|
await service.get_user_state(app_name='my_app', user_id='u1')
|
|
|
|
|
|
def test_database_session_service_visible_in_module_namespace():
|
|
"""DatabaseSessionService must be in dir() so Sphinx autodoc renders it.
|
|
|
|
It is imported lazily via module __getattr__, so without an explicit
|
|
__dir__ it drops out of the generated API reference (issue #4331).
|
|
"""
|
|
import google.adk.sessions as sessions_module
|
|
|
|
assert 'DatabaseSessionService' in dir(sessions_module)
|
|
assert sessions_module.DatabaseSessionService is DatabaseSessionService
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_database_session_service_with_db_url():
|
|
"""Test DatabaseSessionService initialization with db_url."""
|
|
# Test db_url as positional argument
|
|
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
|
|
app_name = 'test_app'
|
|
user_id = 'test_user'
|
|
|
|
# Create and retrieve a session
|
|
session = await service.create_session(
|
|
app_name=app_name, user_id=user_id, state={'key': 'value'}
|
|
)
|
|
assert session.app_name == app_name
|
|
assert session.user_id == user_id
|
|
assert session.state == {'key': 'value'}
|
|
|
|
# Let's check that we can retrieve it
|
|
retrieved = await service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
assert retrieved == session
|
|
|
|
# test db_url as keyword argument
|
|
service2 = DatabaseSessionService(db_url='sqlite+aiosqlite:///:memory:')
|
|
session2 = await service2.create_session(
|
|
app_name=app_name, user_id=user_id, state={'key': 'value2'}
|
|
)
|
|
assert session2.state == {'key': 'value2'}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_database_session_service_with_db_engine():
|
|
"""Test DatabaseSessionService initialization with db_engine."""
|
|
# Create an engine manually with StaticPool to avoid flakes
|
|
engine = create_async_engine(
|
|
'sqlite+aiosqlite:///:memory:',
|
|
poolclass=StaticPool,
|
|
connect_args={'check_same_thread': False},
|
|
)
|
|
|
|
# Create service with db_engine
|
|
service = DatabaseSessionService(db_engine=engine)
|
|
app_name = 'test_app'
|
|
user_id = 'test_user'
|
|
|
|
# Create and retrieve a session
|
|
session = await service.create_session(
|
|
app_name=app_name, user_id=user_id, state={'key': 'value'}
|
|
)
|
|
assert session.app_name == app_name
|
|
assert session.user_id == user_id
|
|
assert session.state == {'key': 'value'}
|
|
|
|
# Let's check that we can retrieve it
|
|
retrieved = await service.get_session(
|
|
app_name=app_name, user_id=user_id, session_id=session.id
|
|
)
|
|
assert retrieved == session
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_database_session_service_caller_owned_engine_not_disposed_on_close():
|
|
"""Verifies that a caller-owned engine is not disposed when the service is closed."""
|
|
engine = create_async_engine(
|
|
'sqlite+aiosqlite:///:memory:',
|
|
poolclass=StaticPool,
|
|
connect_args={'check_same_thread': False},
|
|
)
|
|
|
|
service = DatabaseSessionService(db_engine=engine)
|
|
|
|
# Use the service
|
|
session = await service.create_session(app_name='app', user_id='user')
|
|
assert session is not None
|
|
|
|
# Close the service
|
|
await service.close()
|
|
|
|
# Verify engine is still usable by running a query
|
|
async with engine.connect() as conn:
|
|
result = await conn.execute(text('SELECT 1;'))
|
|
assert result.scalar() == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_database_session_service_requires_one_argument():
|
|
"""Test that DatabaseSessionService requires exactly one of db_url or db_engine."""
|
|
# Neither argument provided
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="Exactly one of 'db_url' or 'db_engine' must be provided",
|
|
):
|
|
DatabaseSessionService()
|
|
|
|
# Both arguments provided
|
|
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="Exactly one of 'db_url' or 'db_engine' must be provided",
|
|
):
|
|
DatabaseSessionService(
|
|
db_url='sqlite+aiosqlite:///:memory:', db_engine=engine
|
|
)
|