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 lines
66 KiB
Python
2055 lines
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
|
||
)
|