From 4ea43933440597156b5ea1d25e42079f72772714 Mon Sep 17 00:00:00 2001
From: John O'Keefe <jrokeefe@protonmail.com>
Date: Fri, 6 Mar 2026 15:03:21 -0500
Subject: [PATCH] refactor(tests): clean up websocket test helper and fix
 broadcast test

- Remove createTestMediaItem helper function and replace with createTestMediaItemID
- Update TestWebSocketProgressBroadcast to use simplified helper
- Add read deadline and initial message read in TestWebSocketUserScopedBroadcast to properly consume initial connection messages
- This reduces code duplication and improves test reliability by properly handling WebSocket connection setup
---
 cmd/server/tests/websocket_test.go | 33 +++++++-----------------------
 internal/app/app.go                | 23 +++------------------
 2 files changed, 10 insertions(+), 46 deletions(-)

diff --git a/cmd/server/tests/websocket_test.go b/cmd/server/tests/websocket_test.go
index b75c1f5..6038725 100644
--- a/cmd/server/tests/websocket_test.go
+++ b/cmd/server/tests/websocket_test.go
@@ -104,8 +104,7 @@ func TestWebSocketProgressBroadcast(t *testing.T) {
 	token := setup.Token
 
 	// Create a test media item
-	userID := getTestUserID(t, setup.DB)
-	mediaID := createTestMediaItem(t, setup.DB, userID)
+	mediaID := createTestMediaItemID(t, setup)
 
 	// Connect WebSocket client
 	wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
@@ -233,30 +232,6 @@ func TestWebSocketInvalidToken(t *testing.T) {
 	assert.Error(t, err)
 }
 
-// Helper function to create a test media item
-func createTestMediaItem(t *testing.T, db *database.Queries, userID uuid.UUID) string {
-	// First create a test library
-	libID, err := db.CreateLibrary(context.Background(), database.CreateLibraryParams{
-		Name:             "Test Library",
-		LibraryTypeID:    pgtype.UUID{Bytes: [16]byte(uuid.UUID{}), Valid: true},
-		CreatedByAdminID: pgtype.UUID{Bytes: [16]byte(userID), Valid: true},
-	})
-	require.NoError(t, err)
-
-	// Create a test media item
-	mediaID, err := db.CreateMediaItem(context.Background(), database.CreateMediaItemParams{
-		LibraryID:      libID.ID,
-		Title:          "Test Book",
-		FilePath:       "/tmp/test.epub",
-		FileSize:       pgtype.Int8{Int64: 1024, Valid: true},
-		MimeType:       pgtype.Text{String: "application/epub+zip", Valid: true},
-		AddedByAdminID: pgtype.UUID{Bytes: [16]byte(userID), Valid: true},
-	})
-	require.NoError(t, err)
-
-	return uuid.UUID(mediaID.ID.Bytes).String()
-}
-
 // TestWebSocketUserScopedBroadcast tests that broadcasts only go to the user who made changes
 func TestWebSocketUserScopedBroadcast(t *testing.T) {
 	setup := setupTestServer(t)
@@ -292,10 +267,16 @@ func TestWebSocketUserScopedBroadcast(t *testing.T) {
 	wsAdmin := connectWebSocketToServer(t, setup.Server.URL, setup.Token)
 	defer wsAdmin.Close()
 
+	wsAdmin.SetReadDeadline(time.Now().Add(2 * time.Second))
+	_, _, _ = wsAdmin.ReadMessage()
+
 	// Connect regular user via WebSocket
 	wsRegular := connectWebSocketToServer(t, setup.Server.URL, setup.RegularToken)
 	defer wsRegular.Close()
 
+	wsRegular.SetReadDeadline(time.Now().Add(2 * time.Second))
+	_, _, _ = wsRegular.ReadMessage()
+
 	// Admin adds book to collection
 	addReq := map[string]interface{}{
 		"book_ids": []string{bookID},
diff --git a/internal/app/app.go b/internal/app/app.go
index 2145c10..bfa7376 100644
--- a/internal/app/app.go
+++ b/internal/app/app.go
@@ -92,32 +92,15 @@ func (a *App) Shutdown() error {
 	// Create context with timeout
 	ctx, cancel := context.WithTimeout(context.Background(), a.shutdownTimeout)
 	defer cancel()
-
-	// Channel to track shutdown completion
-	done := make(chan struct{})
-
 	log.Println("Stopping HTTP Server...")
 	if a.server != nil {
-		ctx, cancel := context.WithTimeout(context.Background(), a.shutdownTimeout)
-		defer cancel()
-
 		if err := a.server.Shutdown(ctx); err != nil {
 			log.Printf("Error stopping HTTP server: %v", err)
 		}
 	}
-
-	log.Println("All services stopped")
-	// Wait for shutdown or timeout
-	select {
-	case <-done:
-		close(a.shutdownDone)
-		log.Println("Graceful shutdown completed successfully")
-		return nil
-	case <-ctx.Done():
-		close(a.shutdownDone)
-		log.Printf("Shutdown timed out after %v", a.shutdownTimeout)
-		return ctx.Err()
-	}
+	close(a.shutdownDone)
+	log.Println("Graceful shutdown completed successfully")
+	return nil
 }
 
 // SetShutdownTimeout sets the maximum time to wait for graceful shutdown