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
+            )
 
 
 # ------------------------------------------------------------------