diff --git a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java index 54f91a858..9be0d11c3 100644 --- a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java +++ b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java @@ -26,6 +26,7 @@ import com.google.cloud.firestore.CollectionReference; import com.google.cloud.firestore.DocumentReference; import com.google.cloud.firestore.DocumentSnapshot; +import com.google.cloud.firestore.FieldValue; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.Query; import com.google.cloud.firestore.QueryDocumentSnapshot; @@ -41,6 +42,7 @@ import java.time.Instant; import java.time.InstantSource; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -132,60 +134,63 @@ public Single createSession( @Nullable Map state, @Nullable String sessionId) { return Single.fromCallable( - () -> { - Objects.requireNonNull(appName, "appName cannot be null"); - Objects.requireNonNull(userId, "userId cannot be null"); + () -> { + Objects.requireNonNull(appName, "appName cannot be null"); + Objects.requireNonNull(userId, "userId cannot be null"); + + String resolvedSessionId = + Optional.ofNullable(sessionId) + .map(String::trim) + .filter(s -> !s.isEmpty()) + .orElseGet(uuidProvider::newUuid); + + ConcurrentMap initialState = + (state == null) ? new ConcurrentHashMap<>() : new ConcurrentHashMap<>(state); + logger.info( + "Creating session for userId: {} with sessionId: {} and initial state: {}", + userId, + resolvedSessionId, + initialState); + List initialEvents = new ArrayList<>(); + Instant now = instantSource.instant(); + Session newSession = + Session.builder(resolvedSessionId) + .appName(appName) + .userId(userId) + .state(initialState) + .events(initialEvents) + .lastUpdateTime(now) + .build(); - String resolvedSessionId = - Optional.ofNullable(sessionId) - .map(String::trim) - .filter(s -> !s.isEmpty()) - .orElseGet(uuidProvider::newUuid); - - ConcurrentMap initialState = - (state == null) ? new ConcurrentHashMap<>() : new ConcurrentHashMap<>(state); - logger.info( - "Creating session for userId: {} with sessionId: {} and initial state: {}", - userId, - resolvedSessionId, - initialState); - List initialEvents = new ArrayList<>(); - Instant now = instantSource.instant(); - Session newSession = - Session.builder(resolvedSessionId) - .appName(appName) - .userId(userId) - .state(initialState) - .events(initialEvents) - .lastUpdateTime(now) - .build(); - - // Convert Session to a Map for Firestore - Map sessionData = new HashMap<>(); - sessionData.put(ID_KEY, newSession.id()); - sessionData.put(APP_NAME_KEY, newSession.appName()); - sessionData.put(USER_ID_KEY, newSession.userId()); - sessionData.put(UPDATE_TIME_KEY, newSession.lastUpdateTime().toString()); - sessionData.put(STATE_KEY, newSession.state()); - String createToken = UUID.randomUUID().toString(); - sessionData.put(CREATE_TOKEN_KEY, createToken); - - // Unlike set(), create() fails if the session already exists instead of replacing it. - DocumentReference sessionDoc = getSessionsCollection(userId).document(resolvedSessionId); - try { - sessionDoc.create(sessionData).get(); - } catch (ExecutionException e) { - if (!(e.getCause() instanceof AlreadyExistsException)) { - throw e; - } - // A retry after a lost reply fails on this call's own write; its token means success. - if (!createToken.equals(sessionDoc.get().get().getString(CREATE_TOKEN_KEY))) { - throw new SessionException(SessionException.SESSION_ALREADY_EXISTS, e.getCause()); - } - } + // Convert Session to a Map for Firestore + Map sessionData = new HashMap<>(); + sessionData.put(ID_KEY, newSession.id()); + sessionData.put(APP_NAME_KEY, newSession.appName()); + sessionData.put(USER_ID_KEY, newSession.userId()); + sessionData.put(UPDATE_TIME_KEY, newSession.lastUpdateTime().toString()); + sessionData.put(STATE_KEY, newSession.state()); + String createToken = UUID.randomUUID().toString(); + sessionData.put(CREATE_TOKEN_KEY, createToken); + + // Unlike set(), create() fails if the session already exists instead of replacing it. + DocumentReference sessionDoc = + getSessionsCollection(userId).document(resolvedSessionId); + try { + sessionDoc.create(sessionData).get(); + } catch (ExecutionException e) { + if (!(e.getCause() instanceof AlreadyExistsException)) { + throw e; + } + // A retry after a lost reply fails on this call's own write; its token means + // success. + if (!createToken.equals(sessionDoc.get().get().getString(CREATE_TOKEN_KEY))) { + throw new SessionException(SessionException.SESSION_ALREADY_EXISTS, e.getCause()); + } + } - return newSession; - }); + return newSession; + }) + .flatMap(session -> mergeWithGlobalState(appName, userId, session)); } /*** @@ -273,6 +278,7 @@ public Maybe getSession( .events(events) .build(); }) + .flatMap(session -> mergeWithGlobalState(appName, userId, session)) .toMaybe(); }); } @@ -491,40 +497,54 @@ private Map eventToMap(Session session, Event event) { @Override public Single listSessions(String appName, String userId) { return Single.fromCallable( - () -> { - Objects.requireNonNull(appName, "appName cannot be null"); - Objects.requireNonNull(userId, "userId cannot be null"); + () -> { + Objects.requireNonNull(appName, "appName cannot be null"); + Objects.requireNonNull(userId, "userId cannot be null"); - logger.info("Listing sessions for userId: {}", userId); + logger.info("Listing sessions for userId: {}", userId); - Query query = getSessionsCollection(userId).whereEqualTo(APP_NAME_KEY, appName); + Query query = getSessionsCollection(userId).whereEqualTo(APP_NAME_KEY, appName); - ApiFuture> querySnapshot = - ApiFutures.transform( - query.get(), // Query is already scoped to the user - snapshot -> snapshot.getDocuments(), - MoreExecutors.directExecutor()); + ApiFuture> querySnapshot = + ApiFutures.transform( + query.get(), // Query is already scoped to the user + snapshot -> snapshot.getDocuments(), + MoreExecutors.directExecutor()); - List sessions = new ArrayList<>(); - for (DocumentSnapshot document : querySnapshot.get()) { - Map data = document.getData(); - if (data != null) { - // Create a session object with empty events and state, as per - // InMemorySessionService - Session session = - Session.builder((String) data.get(ID_KEY)) - .appName((String) data.get(APP_NAME_KEY)) - .userId((String) data.get(USER_ID_KEY)) - .lastUpdateTime(Instant.parse((String) data.get(UPDATE_TIME_KEY))) - .state(new ConcurrentHashMap<>()) // Empty state - .events(new ArrayList<>()) // Empty events - .build(); - sessions.add(session); - } - } + List sessions = new ArrayList<>(); + for (DocumentSnapshot document : querySnapshot.get()) { + Map data = document.getData(); + if (data != null) { + // Create a session object with empty events and state, as per + // InMemorySessionService + Session session = + Session.builder((String) data.get(ID_KEY)) + .appName((String) data.get(APP_NAME_KEY)) + .userId((String) data.get(USER_ID_KEY)) + .lastUpdateTime(Instant.parse((String) data.get(UPDATE_TIME_KEY))) + .state(new ConcurrentHashMap<>()) // Empty state + .events(new ArrayList<>()) // Empty events + .build(); + sessions.add(session); + } + } - return ListSessionsResponse.builder().sessions(sessions).build(); - }); + return sessions; + }) + .flatMap( + sessions -> { + if (sessions.isEmpty()) { + return Single.just(ListSessionsResponse.builder().sessions(sessions).build()); + } + return getGlobalState(appName, userId) + .map( + globalState -> { + for (Session session : sessions) { + session.state().putAll(globalState); + } + return ListSessionsResponse.builder().sessions(sessions).build(); + }); + }); } /** Deletes a session and all its associated events from Firestore. */ @@ -656,14 +676,31 @@ public Single appendEvent(Session session, Event event) { stateDelta.forEach( (key, value) -> { - if (key.startsWith("_app_")) { - appStateUpdates.put(key.substring("_app_".length()), value); - } else if (key.startsWith("_user_")) { - userStateUpdates.put(key.substring("_user_".length()), value); + if (key.startsWith(State.TEMP_PREFIX)) { + // Temp state is ephemeral and not persisted + return; + } else if (key.startsWith(State.APP_PREFIX)) { + String appKey = key.substring(State.APP_PREFIX.length()); + if (value == State.REMOVED) { + appStateUpdates.put(appKey, FieldValue.delete()); + session.state().remove(key); + } else { + appStateUpdates.put(appKey, value); + session.state().put(key, value); + } + } else if (key.startsWith(State.USER_PREFIX)) { + String userKey = key.substring(State.USER_PREFIX.length()); + if (value == State.REMOVED) { + userStateUpdates.put(userKey, FieldValue.delete()); + session.state().remove(key); + } else { + userStateUpdates.put(userKey, value); + session.state().put(key, value); + } } else { // Regular session state sessionStateChanged.set(true); - if (value == null) { + if (value == null || value == State.REMOVED) { session.state().remove(key); } else { session.state().put(key, value); @@ -784,4 +821,66 @@ private T safeCast(Object obj, Class clazz, String fieldName, T defaultVa obj.getClass().getName()); return defaultValue; } + + /** + * Fetches the global application and user state from Firestore and prefixes them with {@link + * State#APP_PREFIX} and {@link State#USER_PREFIX}. + */ + private Single> getGlobalState(String appName, String userId) { + CollectionReference appCol = firestore.collection(APP_STATE_COLLECTION); + DocumentReference appDoc = appCol != null ? appCol.document(appName) : null; + ApiFuture appStateFuture = appDoc != null ? appDoc.get() : null; + + CollectionReference userCol = firestore.collection(USER_STATE_COLLECTION); + DocumentReference userAppDoc = userCol != null ? userCol.document(appName) : null; + CollectionReference usersSubCol = userAppDoc != null ? userAppDoc.collection("users") : null; + DocumentReference userDoc = usersSubCol != null ? usersSubCol.document(userId) : null; + ApiFuture userStateFuture = userDoc != null ? userDoc.get() : null; + + Single> appStateSingle = + appStateFuture != null + ? ApiFutureUtils.toSingle(appStateFuture) + .map( + snap -> { + Map map = new HashMap<>(); + if (snap != null && snap.exists() && snap.getData() != null) { + snap.getData().forEach((k, v) -> map.put(State.APP_PREFIX + k, v)); + } + return map; + }) + : Single.just(Collections.emptyMap()); + + Single> userStateSingle = + userStateFuture != null + ? ApiFutureUtils.toSingle(userStateFuture) + .map( + snap -> { + Map map = new HashMap<>(); + if (snap != null && snap.exists() && snap.getData() != null) { + snap.getData().forEach((k, v) -> map.put(State.USER_PREFIX + k, v)); + } + return map; + }) + : Single.just(Collections.emptyMap()); + + return Single.zip( + appStateSingle, + userStateSingle, + (appMap, userMap) -> { + Map globalState = new HashMap<>(appMap); + globalState.putAll(userMap); + return globalState; + }); + } + + /** Merges app-level and user-level state into the provided session's state map. */ + @CanIgnoreReturnValue + private Single mergeWithGlobalState(String appName, String userId, Session session) { + return getGlobalState(appName, userId) + .map( + globalState -> { + session.state().putAll(globalState); + return session; + }); + } } diff --git a/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java b/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java index eae6de3b4..079ea47a9 100644 --- a/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java +++ b/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java @@ -31,6 +31,7 @@ import com.google.adk.events.Event; import com.google.adk.plugins.BasePlugin; import com.google.adk.sessions.Session; +import com.google.adk.utils.Constants; import com.google.adk.utils.FirestoreProperties; import com.google.api.core.ApiFutures; import com.google.cloud.firestore.CollectionReference; @@ -89,6 +90,8 @@ public void setup() throws Exception { @Mock private DocumentReference mockEventDocRef; @Mock private Query mockQuery; @Mock private QuerySnapshot mockQuerySnapshot; + @Mock private DocumentSnapshot mockAppStateSnapshot; + @Mock private DocumentSnapshot mockUserStateSnapshot; @Mock private WriteResult mockWriteResult; /** @@ -175,6 +178,22 @@ void run_withUserInput_createsSessionAndExecutesAgent() throws Exception { when(mockUserDocRef.collection(anyString())).thenReturn(mockSessionsCollection); when(mockSessionsCollection.document(anyString())).thenReturn(mockSessionDocRef); + // Mock app-state and user-state collections for session global state retrieval + when(mockFirestore.collection(Constants.APP_STATE_COLLECTION)) + .thenReturn(mockAppStateCollection); + when(mockAppStateCollection.document(anyString())).thenReturn(mockAppStateDocRef); + when(mockAppStateDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockAppStateSnapshot)); + when(mockAppStateSnapshot.exists()).thenReturn(false); + + when(mockFirestore.collection(Constants.USER_STATE_COLLECTION)) + .thenReturn(mockUserStateRootCollection); + when(mockUserStateRootCollection.document(anyString())).thenReturn(mockUserStateAppDocRef); + when(mockUserStateAppDocRef.collection(anyString())).thenReturn(mockUserStateUsersCollection); + when(mockUserStateUsersCollection.document(anyString())).thenReturn(mockUserStateUserDocRef); + when(mockUserStateUserDocRef.get()) + .thenReturn(ApiFutures.immediateFuture(mockUserStateSnapshot)); + when(mockUserStateSnapshot.exists()).thenReturn(false); + when(mockSessionDocRef.create(anyMap())) .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); when(mockSessionDocRef.update(anyMap())) diff --git a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java index ba95dc84c..e78d11b5c 100644 --- a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java +++ b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java @@ -38,6 +38,7 @@ import com.google.cloud.firestore.CollectionReference; import com.google.cloud.firestore.DocumentReference; import com.google.cloud.firestore.DocumentSnapshot; +import com.google.cloud.firestore.FieldValue; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.Query; import com.google.cloud.firestore.QueryDocumentSnapshot; @@ -91,6 +92,8 @@ public class FirestoreSessionServiceTest { @Mock private DocumentReference mockUserStateAppDocRef; @Mock private DocumentReference mockUserStateUserDocRef; @Mock private DocumentSnapshot mockSessionSnapshot; + @Mock private DocumentSnapshot mockAppStateSnapshot; + @Mock private DocumentSnapshot mockUserStateSnapshot; @Mock private QueryDocumentSnapshot mockEventSnapshot; @Mock private Query mockQuery; @Mock private QuerySnapshot mockQuerySnapshot; @@ -122,20 +125,28 @@ public void setup() { lenient() .when(mockDb.collection(Constants.APP_STATE_COLLECTION)) .thenReturn(mockAppStateCollection); - lenient().when(mockAppStateCollection.document(APP_NAME)).thenReturn(mockAppStateDocRef); + lenient().when(mockAppStateCollection.document(anyString())).thenReturn(mockAppStateDocRef); + lenient() + .when(mockAppStateDocRef.get()) + .thenReturn(ApiFutures.immediateFuture(mockAppStateSnapshot)); + lenient().when(mockAppStateSnapshot.exists()).thenReturn(false); + lenient() .when(mockDb.collection(Constants.USER_STATE_COLLECTION)) .thenReturn(mockUserStateRootCollection); lenient() - .when(mockUserStateRootCollection.document(APP_NAME)) + .when(mockUserStateRootCollection.document(anyString())) .thenReturn(mockUserStateAppDocRef); lenient() - .when(mockUserStateAppDocRef.collection("users")) + .when(mockUserStateAppDocRef.collection(anyString())) .thenReturn(mockUserStateUsersCollection); lenient() - .when(mockUserStateUsersCollection.document(USER_ID)) + .when(mockUserStateUsersCollection.document(anyString())) .thenReturn(mockUserStateUserDocRef); - + lenient() + .when(mockUserStateUserDocRef.get()) + .thenReturn(ApiFutures.immediateFuture(mockUserStateSnapshot)); + lenient().when(mockUserStateSnapshot.exists()).thenReturn(false); // Default mock for writes lenient() .when(mockSessionDocRef.create(anyMap())) @@ -685,7 +696,7 @@ void appendEvent_withAppOnlyStateDeltas_updatesCorrectStores() { EventActions.builder() .stateDelta( new ConcurrentHashMap<>( - ImmutableMap.of("_app_appKey", "appValue"))) // Corrected type and key + ImmutableMap.of("app:appKey", "appValue"))) // Corrected type and key .build(); Event event = @@ -732,7 +743,7 @@ void appendEvent_withUserOnlyStateDeltas_updatesCorrectStores() { EventActions actions = EventActions.builder() - .stateDelta(new ConcurrentHashMap<>(ImmutableMap.of("_user_userKey", "userValue"))) + .stateDelta(new ConcurrentHashMap<>(ImmutableMap.of("user:userKey", "userValue"))) .build(); Event event = @@ -764,6 +775,116 @@ void appendEvent_withUserOnlyStateDeltas_updatesCorrectStores() { verify(mockSessionDocRef, never()).update(eq(Constants.KEY_STATE), any()); } + /** Tests that appendEvent ignores temp: state keys and never writes them to any store. */ + @Test + void appendEvent_withTempOnlyStateDeltas_doesNotPersistTempState() { + Session session = + Session.builder(SESSION_ID) + .appName(APP_NAME) + .userId(USER_ID) + .state(new ConcurrentHashMap<>()) + .build(); + + EventActions actions = + EventActions.builder() + .stateDelta( + new ConcurrentHashMap<>( + ImmutableMap.of( + State.TEMP_PREFIX + "scratchKey", + "tempValue", + State.TEMP_PREFIX + "counter", + 10))) + .build(); + + Event event = + Event.builder() + .author("model") + .content(Content.builder().parts(List.of(Part.fromText("..."))).build()) + .actions(actions) + .build(); + + // Mock Firestore interactions for appendEvent + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockEventsCollection.document()).thenReturn(mockEventDocRef); + when(mockEventDocRef.getId()).thenReturn(EVENT_ID); + when(mockEventsCollection.document(EVENT_ID)).thenReturn(mockEventDocRef); + when(mockSessionDocRef.update(anyMap())) + .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); + + // Act + sessionService.appendEvent(session, event).test().assertComplete(); + + // Assert + // Ephemeral temp keys must never be added to session state + assertThat(session.state()).doesNotContainKey(State.TEMP_PREFIX + "scratchKey"); + assertThat(session.state()).doesNotContainKey(State.TEMP_PREFIX + "counter"); + + // Scoped stores and session state update must never be invoked + verify(mockAppStateDocRef, never()).set(anyMap(), any(SetOptions.class)); + verify(mockUserStateUsersCollection, never()).document(anyString()); + verify(mockSessionDocRef, never()).update(eq(Constants.KEY_STATE), any()); + } + + /** Tests that appendEvent with State.REMOVED deletes keys from app, user, and session stores. */ + @Test + @SuppressWarnings("unchecked") + void appendEvent_withStateRemovedSentinel_deletesKeysAcrossStores() { + // Arrange + ConcurrentHashMap sessionState = new ConcurrentHashMap<>(); + sessionState.put("sessionKey", "initialValue"); + + Session session = + Session.builder(SESSION_ID).appName(APP_NAME).userId(USER_ID).state(sessionState).build(); + + EventActions actions = + EventActions.builder() + .stateDelta( + new ConcurrentHashMap<>( + ImmutableMap.of( + State.APP_PREFIX + "appKey", + State.REMOVED, + State.USER_PREFIX + "userKey", + State.REMOVED, + "sessionKey", + State.REMOVED))) + .build(); + + Event event = + Event.builder() + .author("model") + .content(Content.builder().parts(List.of(Part.fromText("..."))).build()) + .actions(actions) + .build(); + + // Mock Firestore interactions for appendEvent + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockEventsCollection.document()).thenReturn(mockEventDocRef); + when(mockEventDocRef.getId()).thenReturn(EVENT_ID); + when(mockEventsCollection.document(EVENT_ID)).thenReturn(mockEventDocRef); + when(mockSessionDocRef.update(anyMap())) + .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); + when(mockSessionDocRef.update(eq(Constants.KEY_STATE), any())) + .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); + + // Act + sessionService.appendEvent(session, event).test().assertComplete(); + + // Assert + // Verify app state uses FieldValue.delete() + ArgumentCaptor> appStateCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockAppStateDocRef).set(appStateCaptor.capture(), any(SetOptions.class)); + assertThat(appStateCaptor.getValue()).containsEntry("appKey", FieldValue.delete()); + + // Verify user state uses FieldValue.delete() + ArgumentCaptor> userStateCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockUserStateUserDocRef).set(userStateCaptor.capture(), any(SetOptions.class)); + assertThat(userStateCaptor.getValue()).containsEntry("userKey", FieldValue.delete()); + + // Verify session state entry was removed from in-memory session object + assertThat(session.state()).doesNotContainKey("sessionKey"); + verify(mockSessionDocRef).update(eq(Constants.KEY_STATE), any()); + } + /** Tests that getSession skips malformed events and returns only the well-formed ones. */ @Test @SuppressWarnings("unchecked") @@ -1122,6 +1243,105 @@ void deleteSession_appNameMismatch_doesNotDelete() { verify(mockWriteBatch, never()).commit(); } + /** Tests that getSession merges app and user state with proper prefixes. */ + @Test + void getSession_mergesGlobalAppStateAndUserState() { + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockSessionDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockSessionSnapshot)); + when(mockSessionSnapshot.exists()).thenReturn(true); + when(mockSessionSnapshot.getData()) + .thenReturn( + ImmutableMap.of( + Constants.KEY_ID, SESSION_ID, + Constants.KEY_APP_NAME, APP_NAME, + Constants.KEY_USER_ID, USER_ID, + Constants.KEY_UPDATE_TIME, NOW.toString(), + Constants.KEY_STATE, ImmutableMap.of("sessionKey", "sessionValue"))); + when(mockSessionSnapshot.getReference()).thenReturn(mockSessionDocRef); + when(mockQuery.get()).thenReturn(ApiFutures.immediateFuture(mockQuerySnapshot)); + when(mockQuerySnapshot.getDocuments()).thenReturn(Collections.emptyList()); + + when(mockAppStateSnapshot.exists()).thenReturn(true); + when(mockAppStateSnapshot.getData()).thenReturn(ImmutableMap.of("theme", "dark")); + when(mockUserStateSnapshot.exists()).thenReturn(true); + when(mockUserStateSnapshot.getData()).thenReturn(ImmutableMap.of("name", "Jae")); + + TestObserver testObserver = + sessionService.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.empty()).test(); + + testObserver.assertComplete(); + testObserver.assertValue( + s -> { + assertThat(s.state()).containsEntry("sessionKey", "sessionValue"); + assertThat(s.state()).containsEntry("app:theme", "dark"); + assertThat(s.state()).containsEntry("user:name", "Jae"); + return true; + }); + } + + /** Tests that createSession merges existing global app and user state into the new session. */ + @Test + void createSession_mergesGlobalAppStateAndUserState() { + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockAppStateSnapshot.exists()).thenReturn(true); + when(mockAppStateSnapshot.getData()).thenReturn(ImmutableMap.of("theme", "dark")); + when(mockUserStateSnapshot.exists()).thenReturn(true); + when(mockUserStateSnapshot.getData()).thenReturn(ImmutableMap.of("name", "Jae")); + + TestObserver testObserver = + sessionService + .createSession( + APP_NAME, + USER_ID, + new ConcurrentHashMap<>(ImmutableMap.of("initial", "val")), + SESSION_ID) + .test(); + + testObserver.assertComplete(); + testObserver.assertValue( + s -> { + assertThat(s.state()).containsEntry("initial", "val"); + assertThat(s.state()).containsEntry("app:theme", "dark"); + assertThat(s.state()).containsEntry("user:name", "Jae"); + return true; + }); + } + + /** Tests that listSessions merges existing global app and user state into listed sessions. */ + @Test + void listSessions_mergesGlobalAppStateAndUserState() { + when(mockSessionsCollection.whereEqualTo(Constants.KEY_APP_NAME, APP_NAME)) + .thenReturn(mockQuery); + when(mockQuery.get()).thenReturn(ApiFutures.immediateFuture(mockQuerySnapshot)); + QueryDocumentSnapshot mockDoc = mock(QueryDocumentSnapshot.class); + when(mockQuerySnapshot.getDocuments()).thenReturn(List.of(mockDoc)); + when(mockDoc.getData()) + .thenReturn( + ImmutableMap.of( + Constants.KEY_ID, SESSION_ID, + Constants.KEY_APP_NAME, APP_NAME, + Constants.KEY_USER_ID, USER_ID, + Constants.KEY_UPDATE_TIME, NOW.toString())); + + when(mockAppStateSnapshot.exists()).thenReturn(true); + when(mockAppStateSnapshot.getData()).thenReturn(ImmutableMap.of("theme", "dark")); + when(mockUserStateSnapshot.exists()).thenReturn(true); + when(mockUserStateSnapshot.getData()).thenReturn(ImmutableMap.of("name", "Jae")); + + TestObserver testObserver = + sessionService.listSessions(APP_NAME, USER_ID).test(); + + testObserver.assertComplete(); + testObserver.assertValue( + response -> { + assertThat(response.sessions()).hasSize(1); + Session s = response.sessions().get(0); + assertThat(s.state()).containsEntry("app:theme", "dark"); + assertThat(s.state()).containsEntry("user:name", "Jae"); + return true; + }); + } + private static AlreadyExistsException alreadyExists() { return new AlreadyExistsException( "Document already exists",