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
This commit is contained in:
@@ -104,8 +104,7 @@ func TestWebSocketProgressBroadcast(t *testing.T) {
|
|||||||
token := setup.Token
|
token := setup.Token
|
||||||
|
|
||||||
// Create a test media item
|
// Create a test media item
|
||||||
userID := getTestUserID(t, setup.DB)
|
mediaID := createTestMediaItemID(t, setup)
|
||||||
mediaID := createTestMediaItem(t, setup.DB, userID)
|
|
||||||
|
|
||||||
// Connect WebSocket client
|
// Connect WebSocket client
|
||||||
wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
|
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)
|
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
|
// TestWebSocketUserScopedBroadcast tests that broadcasts only go to the user who made changes
|
||||||
func TestWebSocketUserScopedBroadcast(t *testing.T) {
|
func TestWebSocketUserScopedBroadcast(t *testing.T) {
|
||||||
setup := setupTestServer(t)
|
setup := setupTestServer(t)
|
||||||
@@ -292,10 +267,16 @@ func TestWebSocketUserScopedBroadcast(t *testing.T) {
|
|||||||
wsAdmin := connectWebSocketToServer(t, setup.Server.URL, setup.Token)
|
wsAdmin := connectWebSocketToServer(t, setup.Server.URL, setup.Token)
|
||||||
defer wsAdmin.Close()
|
defer wsAdmin.Close()
|
||||||
|
|
||||||
|
wsAdmin.SetReadDeadline(time.Now().Add(2 * time.Second))
|
||||||
|
_, _, _ = wsAdmin.ReadMessage()
|
||||||
|
|
||||||
// Connect regular user via WebSocket
|
// Connect regular user via WebSocket
|
||||||
wsRegular := connectWebSocketToServer(t, setup.Server.URL, setup.RegularToken)
|
wsRegular := connectWebSocketToServer(t, setup.Server.URL, setup.RegularToken)
|
||||||
defer wsRegular.Close()
|
defer wsRegular.Close()
|
||||||
|
|
||||||
|
wsRegular.SetReadDeadline(time.Now().Add(2 * time.Second))
|
||||||
|
_, _, _ = wsRegular.ReadMessage()
|
||||||
|
|
||||||
// Admin adds book to collection
|
// Admin adds book to collection
|
||||||
addReq := map[string]interface{}{
|
addReq := map[string]interface{}{
|
||||||
"book_ids": []string{bookID},
|
"book_ids": []string{bookID},
|
||||||
|
|||||||
+3
-20
@@ -92,32 +92,15 @@ func (a *App) Shutdown() error {
|
|||||||
// Create context with timeout
|
// Create context with timeout
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), a.shutdownTimeout)
|
ctx, cancel := context.WithTimeout(context.Background(), a.shutdownTimeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
// Channel to track shutdown completion
|
|
||||||
done := make(chan struct{})
|
|
||||||
|
|
||||||
log.Println("Stopping HTTP Server...")
|
log.Println("Stopping HTTP Server...")
|
||||||
if a.server != nil {
|
if a.server != nil {
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), a.shutdownTimeout)
|
|
||||||
defer cancel()
|
|
||||||
|
|
||||||
if err := a.server.Shutdown(ctx); err != nil {
|
if err := a.server.Shutdown(ctx); err != nil {
|
||||||
log.Printf("Error stopping HTTP server: %v", err)
|
log.Printf("Error stopping HTTP server: %v", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
close(a.shutdownDone)
|
||||||
log.Println("All services stopped")
|
log.Println("Graceful shutdown completed successfully")
|
||||||
// Wait for shutdown or timeout
|
return nil
|
||||||
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()
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// SetShutdownTimeout sets the maximum time to wait for graceful shutdown
|
// SetShutdownTimeout sets the maximum time to wait for graceful shutdown
|
||||||
|
|||||||
Reference in New Issue
Block a user