From be847a38abc82adbc45278928374496ae7a32506 Mon Sep 17 00:00:00 2001 From: A8080816 Date: Mon, 23 Feb 2026 19:04:13 +0000 Subject: [PATCH] test: refactor test --- tests/test_compaction.py | 58 +++++++++++++++++++++++++++++++--------- 1 file changed, 45 insertions(+), 13 deletions(-) diff --git a/tests/test_compaction.py b/tests/test_compaction.py index 3e7b232..eca75a2 100644 --- a/tests/test_compaction.py +++ b/tests/test_compaction.py @@ -14,7 +14,8 @@ from google.adk.events.event import Event from google.cloud.firestore_v1.async_client import AsyncClient from google.genai.types import Content, GenerateContentResponseUsageMetadata, Part -from va_agent.session import FirestoreSessionService, _try_claim_compaction_txn +from va_agent.session import FirestoreSessionService +from va_agent.compaction import SessionCompactor, _try_claim_compaction_txn pytestmark = pytest.mark.asyncio @@ -178,7 +179,9 @@ class TestCompactionEdgeCases: await compaction_service.append_event(session, e) # Trigger compaction manually even though threshold wouldn't fire - await compaction_service._compact_session(session) + events_ref = compaction_service._events_col(app_name, user_id, session.id) + session_ref = compaction_service._session_ref(app_name, user_id, session.id) + await compaction_service._compactor._compact_session(session, events_ref, session_ref) mock_genai_client.aio.models.generate_content.assert_not_called() @@ -205,7 +208,9 @@ class TestCompactionEdgeCases: ) # Should not raise - await compaction_service._compact_session(session) + events_ref = compaction_service._events_col(app_name, user_id, session.id) + session_ref = compaction_service._session_ref(app_name, user_id, session.id) + await compaction_service._compactor._compact_session(session, events_ref, session_ref) # All events should still be present fetched = await compaction_service.get_session( @@ -268,7 +273,7 @@ class TestEventsToText: invocation_id="inv-2", ), ] - text = FirestoreSessionService._events_to_text(events) + text = SessionCompactor._events_to_text(events) assert "User: Hi there" in text assert "Assistant: Hello!" in text @@ -280,7 +285,7 @@ class TestEventsToText: invocation_id="inv-1", ), ] - text = FirestoreSessionService._events_to_text(events) + text = SessionCompactor._events_to_text(events) assert text == "" @@ -368,11 +373,15 @@ class TestGuardedCompact: # Hold the in-process lock so _guarded_compact skips key = f"{app_name}__{user_id}__{session.id}" - lock = compaction_service._compaction_locks.setdefault( + lock = compaction_service._compactor._compaction_locks.setdefault( key, asyncio.Lock() ) + events_ref = compaction_service._events_col(app_name, user_id, session.id) + session_ref = compaction_service._session_ref(app_name, user_id, session.id) async with lock: - await compaction_service._guarded_compact(session) + await compaction_service._compactor.guarded_compact( + session, events_ref, session_ref + ) mock_genai_client.aio.models.generate_content.assert_not_called() @@ -399,7 +408,10 @@ class TestGuardedCompact: ) await session_ref.update({"compaction_lock": time.time()}) - await compaction_service._guarded_compact(session) + events_ref = compaction_service._events_col(app_name, user_id, session.id) + await compaction_service._compactor.guarded_compact( + session, events_ref, session_ref + ) mock_genai_client.aio.models.generate_content.assert_not_called() @@ -411,10 +423,18 @@ class TestGuardedCompact: ) with patch( - "va_agent.session._try_claim_compaction_txn", + "va_agent.compaction._try_claim_compaction_txn", side_effect=RuntimeError("Firestore down"), ): - await compaction_service._guarded_compact(session) + events_ref = compaction_service._events_col( + app_name, user_id, session.id + ) + session_ref = compaction_service._session_ref( + app_name, user_id, session.id + ) + await compaction_service._compactor.guarded_compact( + session, events_ref, session_ref + ) mock_genai_client.aio.models.generate_content.assert_not_called() @@ -427,11 +447,19 @@ class TestGuardedCompact: # Make _compact_session raise an unhandled exception with patch.object( - compaction_service, + compaction_service._compactor, "_compact_session", side_effect=RuntimeError("unexpected crash"), ): - await compaction_service._guarded_compact(session) + events_ref = compaction_service._events_col( + app_name, user_id, session.id + ) + session_ref = compaction_service._session_ref( + app_name, user_id, session.id + ) + await compaction_service._compactor.guarded_compact( + session, events_ref, session_ref + ) # Lock should be released even after failure session_ref = compaction_service._session_ref( @@ -467,7 +495,11 @@ class TestGuardedCompact: side_effect=patched_session_ref, ): # Should not raise despite lock release failure - await compaction_service._guarded_compact(session) + events_ref = compaction_service._events_col(app_name, user_id, session.id) + session_ref = compaction_service._session_ref(app_name, user_id, session.id) + await compaction_service._compactor.guarded_compact( + session, events_ref, session_ref + ) # ------------------------------------------------------------------