SessionManager: Refactor checkSession to throw instead of return bool (…

…#8) Also rename the method accordingly. Signed-off-by: Florian Hotze <dev@florianhotze.com>
llamara-ai · Jan 1, 2025 · c4c73a6 · c4c73a6
1 parent d6fab0d
commit c4c73a6
Show file tree

Hide file tree

Showing 7 changed files with 33 additions and 43 deletions.
diff --git a/src/main/java/com/github/llamara/ai/internal/internal/rest/ChatResource.java b/src/main/java/com/github/llamara/ai/internal/internal/rest/ChatResource.java
@@ -27,7 +27,6 @@
 import jakarta.ws.rs.Consumes;
 import jakarta.ws.rs.DELETE;
 import jakarta.ws.rs.GET;
-import jakarta.ws.rs.NotFoundException;
 import jakarta.ws.rs.POST;
 import jakarta.ws.rs.PUT;
 import jakarta.ws.rs.Path;
@@ -132,9 +131,7 @@ public String chat(
                     UUID sessionId,
             String prompt)
             throws ChatModelNotFoundException, SessionNotFoundException {
-        if (!sessionManager.checkSession(sessionId)) {
-            throw new SessionNotFoundException(sessionId);
-        }
+        sessionManager.enforceSessionValid(sessionId);
         ChatModelAiService chatModelAiService = chatModelProvider.getModel(uid).service();
         return chatModelAiService.chat(sessionId, !identity.isAnonymous(), prompt);
     }
@@ -173,10 +170,8 @@ public Multi<String> streamingPrompt(
                             required = true)
                     UUID sessionId,
             String prompt)
-            throws ChatModelNotFoundException {
-        if (!sessionManager.checkSession(sessionId)) {
-            throw new NotFoundException("Session not found.");
-        }
+            throws ChatModelNotFoundException, SessionNotFoundException {
+        sessionManager.enforceSessionValid(sessionId);
         ChatModelAiService chatModelAiService = chatModelProvider.getModel(uid).service();
         Multi<String> sourceMulti =
                 Multi.createFrom()

diff --git a/...github/llamara/ai/internal/internal/security/session/AnonymousUserSessionManagerImpl.java b/...github/llamara/ai/internal/internal/security/session/AnonymousUserSessionManagerImpl.java
@@ -84,12 +84,11 @@ public void shutdown() {
     }
 
     @Override
-    public boolean checkSession(UUID sessionId) {
-        boolean valid = anonymousSessions.contains(sessionId);
-        if (valid) {
-            scheduleDeletion(sessionId);
+    public void enforceSessionValid(UUID sessionId) throws SessionNotFoundException {
+        if (!anonymousSessions.contains(sessionId)) {
+            throw new SessionNotFoundException(sessionId);
         }
-        return valid;
+        scheduleDeletion(sessionId);
     }
 
     private void scheduleDeletion(UUID sessionId) {

diff --git a/...ub/llamara/ai/internal/internal/security/session/AuthenticatedUserSessionManagerImpl.java b/...ub/llamara/ai/internal/internal/security/session/AuthenticatedUserSessionManagerImpl.java
@@ -71,13 +71,8 @@ public class AuthenticatedUserSessionManagerImpl implements SessionManager {
     }
 
     @Override
-    public boolean checkSession(UUID sessionId) {
-        try {
-            userAwareSessionRepository.findById(sessionId);
-        } catch (SessionNotFoundException e) {
-            return false;
-        }
-        return true;
+    public void enforceSessionValid(UUID sessionId) throws SessionNotFoundException {
+        userAwareSessionRepository.findById(sessionId);
     }
 
     private User getUser() {
@@ -129,9 +124,7 @@ public void deleteSession(UUID sessionId) throws SessionNotFoundException {
     @Override
     public Uni<List<ChatMessageRecord>> getChatHistory(UUID sessionId)
             throws SessionNotFoundException {
-        if (!checkSession(sessionId)) {
-            throw new SessionNotFoundException(sessionId);
-        }
+        enforceSessionValid(sessionId);
         return chatHistoryStore.getMessages(sessionId);
     }
 }
diff --git a/src/main/java/com/github/llamara/ai/internal/internal/security/session/SessionManager.java b/src/main/java/com/github/llamara/ai/internal/internal/security/session/SessionManager.java
@@ -39,12 +39,13 @@
  */
 public interface SessionManager {
     /**
-     * Check whether the given session is valid for the current user.
+     * Enforces that the session with the given ID is valid for the current user. If the session is
+     * not valid, a {@link SessionNotFoundException} is thrown.
      *
      * @param sessionId the ID of the session to check
-     * @return {@code true} if the session is valid, {@code false} otherwise
+     * @throws SessionNotFoundException if no session with the given ID was found for the current
      */
-    boolean checkSession(UUID sessionId);
+    void enforceSessionValid(UUID sessionId) throws SessionNotFoundException;
 
     /**
      * Get all sessions for the current user.

diff --git a/src/main/java/com/github/llamara/ai/internal/internal/security/user/UserManager.java b/src/main/java/com/github/llamara/ai/internal/internal/security/user/UserManager.java
@@ -21,7 +21,7 @@ public interface UserManager {
      * Enforces that the user is registered. If the user is not registered, a {@link
      * UserNotRegisteredException} is thrown.
      *
-     * @throws UserNotRegisteredException
+     * @throws UserNotRegisteredException if the user is not registered
      */
     void enforceRegistered() throws UserNotRegisteredException;
 

diff --git a/...ub/llamara/ai/internal/internal/security/session/AnonymousUserSessionManagerImplTest.java b/...ub/llamara/ai/internal/internal/security/session/AnonymousUserSessionManagerImplTest.java
@@ -27,9 +27,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -90,25 +88,27 @@ void shutdownDeletesAllSessions() {
     }
 
     @Test
-    void checkSessionReturnsFalseForNonExistingSession() {
-        assertFalse(sessionManager.checkSession(UUID.randomUUID()));
+    void enforceSessionValidThrowsSessionNotFoundExceptionForNonExistingSession() {
+        assertThrows(
+                SessionNotFoundException.class,
+                () -> sessionManager.enforceSessionValid(UUID.randomUUID()));
     }
 
     @Test
-    void checkSessionReturnsTrueForExistingSession() {
+    void enforceSessionValidDoesNotThrowSessionNotFoundExceptionForExistingSession() {
         UUID sessionId = sessionManager.createSession().getId();
-        assertTrue(sessionManager.checkSession(sessionId));
+        assertDoesNotThrow(() -> sessionManager.enforceSessionValid(sessionId));
     }
 
     @Test
-    void checkSessionPostponesDeletionForExistingSession() throws InterruptedException {
+    void enforceSessionValidPostponesDeletionForExistingSession() throws InterruptedException {
         when(securityConfig.anonymousUserSessionTimeout()).thenReturn(1);
 
         UUID sessionId = sessionManager.createSession().getId();
         verify(chatMemoryStore, never()).deleteMessages(sessionId);
 
         when(securityConfig.anonymousUserSessionTimeout()).thenReturn(2);
-        sessionManager.checkSession(sessionId);
+        assertDoesNotThrow(() -> sessionManager.enforceSessionValid(sessionId));
         verify(chatMemoryStore, never()).deleteMessages(sessionId);
         Thread.sleep(2500); // NOSONAR: sleep is necessary to wait for scheduled deletion
         verify(chatMemoryStore, times(1)).deleteMessages(sessionId);

diff --git a/...lamara/ai/internal/internal/security/session/AuthenticatedUserSessionManagerImplTest.java b/...lamara/ai/internal/internal/security/session/AuthenticatedUserSessionManagerImplTest.java
@@ -27,9 +27,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.*;
@@ -109,8 +107,10 @@ void setup() {
         }
 
         @Test
-        void checkSessionReturnsFalse() {
-            assertFalse(sessionManager.checkSession(UUID.randomUUID()));
+        void enforceSessionValidThrowsSessionNotFoundException() {
+            assertThrows(
+                    SessionNotFoundException.class,
+                    () -> sessionManager.enforceSessionValid(UUID.randomUUID()));
         }
 
         @Test
@@ -157,8 +157,10 @@ void setup() {
         }
 
         @Test
-        void checkSessionReturnsFalseForForeignSession() {
-            assertFalse(sessionManager.checkSession(foreignSessionId));
+        void enforceSessionValidThrowsSessionNotFoundExceptionForForeignSession() {
+            assertThrows(
+                    SessionNotFoundException.class,
+                    () -> sessionManager.enforceSessionValid(foreignSessionId));
         }
 
         @Test
@@ -194,8 +196,8 @@ void setup() {
         }
 
         @Test
-        void checkSessionReturnsTrueForOwnSession() {
-            assertTrue(sessionManager.checkSession(ownSessionId));
+        void enforceSessionValidDoesNotThrowSessionNotFoundExceptionForOwnSession() {
+            assertDoesNotThrow(() -> sessionManager.enforceSessionValid(ownSessionId));
         }
 
         @Test