john-okeefe
4ea4393344
- 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
339 lines
11 KiB
Go
339 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"bookhoard/internal/database"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/gorilla/websocket"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// TestWebSocketConnection tests basic WebSocket connection and authentication
|
|
/* func TestWebSocketConnection(t *testing.T) {
|
|
// Setup test server with WebSocket
|
|
setup := setupTestServer(t)
|
|
|
|
// Get JWT token for a test user
|
|
token := setup.Token
|
|
|
|
// Connect to WebSocket endpoint
|
|
wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
|
|
|
|
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
require.NoError(t, err, "Failed to connect to WebSocket")
|
|
defer ws.Close()
|
|
|
|
// Set read deadline
|
|
ws.SetReadDeadline(time.Now().Add(5 * time.Second))
|
|
|
|
// Wait for initial state message
|
|
_, msg, err := ws.ReadMessage()
|
|
require.NoError(t, err, "Failed to read initial message")
|
|
|
|
var initialMsg map[string]interface{}
|
|
err = json.Unmarshal(msg, &initialMsg)
|
|
require.NoError(t, err)
|
|
|
|
assert.Equal(t, "initial_state", initialMsg["type"])
|
|
assert.Contains(t, initialMsg["data"], "progress")
|
|
assert.Contains(t, initialMsg["data"], "devices")
|
|
} */
|
|
|
|
// TestWebSocketDeviceAuth tests device authentication via WebSocket
|
|
func TestWebSocketDeviceAuth(t *testing.T) {
|
|
setup := setupTestServer(t)
|
|
|
|
// Create a test device
|
|
userID := getTestUserID(t, setup.DB)
|
|
deviceID := uuid.New()
|
|
|
|
_, err := setup.DB.CreateDevice(context.Background(), database.CreateDeviceParams{
|
|
UserID: pgtype.UUID{Bytes: [16]byte(userID), Valid: true},
|
|
DeviceName: "Test KOReader",
|
|
DeviceType: "koreader",
|
|
DeviceIdentifier: "test-device-123",
|
|
AuthToken: "test-device-token-" + deviceID.String(),
|
|
SyncEnabled: pgtype.Bool{Bool: true, Valid: true},
|
|
AutoSync: pgtype.Bool{Bool: true, Valid: true},
|
|
SyncFrequencyMinutes: pgtype.Int4{Int32: 5, Valid: true},
|
|
DeviceMetadata: []byte("{}"),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Use gorilla/websocket's RequestHeader support
|
|
wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=device-auth-test"
|
|
// Create dialer with custom headers
|
|
dialer := &websocket.Dialer{
|
|
HandshakeTimeout: 5 * time.Second,
|
|
}
|
|
headers := http.Header{}
|
|
headers.Set("Authorization", "Bearer test-device-token-"+deviceID.String())
|
|
// Connect with device token in header
|
|
ws, resp, err := dialer.Dial(wsURL, headers)
|
|
require.NoError(t, err, "WebSocket connection with device token should succeed")
|
|
defer ws.Close()
|
|
if resp != nil {
|
|
defer resp.Body.Close()
|
|
require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode, "Should upgrade to WebSocket")
|
|
}
|
|
// Read initial state message
|
|
ws.SetReadDeadline(time.Now().Add(5 * time.Second))
|
|
_, msg, err := ws.ReadMessage()
|
|
require.NoError(t, err, "Should receive initial state message")
|
|
var initialMsg map[string]interface{}
|
|
err = json.Unmarshal(msg, &initialMsg)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "initial_state", initialMsg["type"])
|
|
}
|
|
|
|
// TestWebSocketProgressBroadcast tests that progress updates are broadcast to connected clients
|
|
func TestWebSocketProgressBroadcast(t *testing.T) {
|
|
setup := setupTestServer(t)
|
|
|
|
// Get JWT token
|
|
token := setup.Token
|
|
|
|
// Create a test media item
|
|
mediaID := createTestMediaItemID(t, setup)
|
|
|
|
// Connect WebSocket client
|
|
wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
|
|
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
require.NoError(t, err)
|
|
defer ws.Close()
|
|
|
|
// Read and discard initial state message
|
|
ws.SetReadDeadline(time.Now().Add(5 * time.Second))
|
|
_, _, _ = ws.ReadMessage()
|
|
|
|
// Update progress via HTTP API
|
|
progressReq := map[string]interface{}{
|
|
"source": "test",
|
|
"location": map[string]interface{}{
|
|
"percentage": 0.5,
|
|
},
|
|
"device_metadata": map[string]interface{}{
|
|
"device_type": "web",
|
|
},
|
|
}
|
|
body, _ := json.Marshal(progressReq)
|
|
|
|
req, _ := http.NewRequest("POST", setup.Server.URL+"/api/progress/"+mediaID, strings.NewReader(string(body)))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
|
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
require.NoError(t, err)
|
|
defer resp.Body.Close()
|
|
|
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
|
|
|
// Read the broadcast message from WebSocket
|
|
_, msg, err := ws.ReadMessage()
|
|
require.NoError(t, err, "Failed to read broadcast message")
|
|
|
|
var broadcastMsg map[string]interface{}
|
|
err = json.Unmarshal(msg, &broadcastMsg)
|
|
require.NoError(t, err)
|
|
|
|
assert.Equal(t, "progress_update", broadcastMsg["type"])
|
|
|
|
data := broadcastMsg["data"].(map[string]interface{})
|
|
assert.Equal(t, mediaID, data["book_id"])
|
|
assert.InDelta(t, 0.5, data["percentage"], 0.01)
|
|
|
|
sourceDevice := broadcastMsg["source_device"].(map[string]interface{})
|
|
assert.Equal(t, "web", sourceDevice["type"])
|
|
}
|
|
|
|
// TestWebSocketPingPong tests that ping/pong messages work correctly
|
|
func TestWebSocketPingPong(t *testing.T) {
|
|
setup := setupTestServer(t)
|
|
|
|
token := setup.Token
|
|
|
|
wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
|
|
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
require.NoError(t, err)
|
|
defer ws.Close()
|
|
|
|
// Read initial message
|
|
ws.SetReadDeadline(time.Now().Add(5 * time.Second))
|
|
_, _, _ = ws.ReadMessage()
|
|
|
|
// Send a ping message (as a text message for testing)
|
|
err = ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"ping"}`))
|
|
require.NoError(t, err)
|
|
|
|
// Server should respond with pong
|
|
ws.SetReadDeadline(time.Now().Add(2 * time.Second))
|
|
_, msg, err := ws.ReadMessage()
|
|
if err == nil {
|
|
var pongMsg map[string]interface{}
|
|
err = json.Unmarshal(msg, &pongMsg)
|
|
if err == nil {
|
|
// Server might respond with pong
|
|
assert.Equal(t, "pong", pongMsg["type"])
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestWebSocketConnectionLimit tests that the server handles multiple connections
|
|
func TestWebSocketConnectionLimit(t *testing.T) {
|
|
setup := setupTestServer(t)
|
|
|
|
token := setup.Token
|
|
|
|
// Create multiple connections
|
|
connections := make([]*websocket.Conn, 5)
|
|
for i := 0; i < 5; i++ {
|
|
wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + fmt.Sprintf("/ws/sync?token=%s&conn=%d", token, i)
|
|
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
require.NoError(t, err, "Failed to create connection %d", i)
|
|
connections[i] = ws
|
|
|
|
// Read initial message
|
|
ws.SetReadDeadline(time.Now().Add(5 * time.Second))
|
|
_, _, _ = ws.ReadMessage()
|
|
}
|
|
|
|
// Close all connections
|
|
for _, ws := range connections {
|
|
ws.Close()
|
|
}
|
|
}
|
|
|
|
// TestWebSocketInvalidToken tests that invalid tokens are rejected
|
|
func TestWebSocketInvalidToken(t *testing.T) {
|
|
setup := setupTestServer(t)
|
|
|
|
// Try to connect with invalid token
|
|
wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=invalid-token"
|
|
|
|
_, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
assert.Error(t, err, "Expected error when connecting with invalid token")
|
|
|
|
// Check if it's a websocket close error
|
|
if websocket.IsCloseError(err, 1000, 1001, 1002, 1003, 1005, 1006, 1007, 1008, 1009, 1010, 1011) {
|
|
// Expected close error
|
|
return
|
|
}
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// TestWebSocketUserScopedBroadcast tests that broadcasts only go to the user who made changes
|
|
func TestWebSocketUserScopedBroadcast(t *testing.T) {
|
|
setup := setupTestServer(t)
|
|
client := &http.Client{}
|
|
|
|
// Use pre-created admin user (setup.Token) and regular user (setup.RegularToken)
|
|
// Both users already created by setupTestServer()
|
|
|
|
// Create a collection for admin user
|
|
collectionReq := map[string]interface{}{
|
|
"name": "Admin Collection",
|
|
"description": "Test collection for WebSocket test",
|
|
}
|
|
collectionBody, _ := json.Marshal(collectionReq)
|
|
|
|
collectionHTTP, _ := http.NewRequest("POST", setup.Server.URL+"/api/collections", bytes.NewBuffer(collectionBody))
|
|
collectionHTTP.Header.Set("Content-Type", "application/json")
|
|
collectionHTTP.Header.Set("Authorization", "Bearer "+setup.Token)
|
|
|
|
collectionResp, err := client.Do(collectionHTTP)
|
|
require.NoError(t, err)
|
|
defer collectionResp.Body.Close()
|
|
require.Equal(t, http.StatusCreated, collectionResp.StatusCode)
|
|
|
|
var collectionResult map[string]interface{}
|
|
json.NewDecoder(collectionResp.Body).Decode(&collectionResult)
|
|
collectionID := collectionResult["id"].(string)
|
|
|
|
// Create a test book via API
|
|
bookID := createTestMediaItemID(t, setup)
|
|
|
|
// Connect admin user via WebSocket
|
|
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},
|
|
}
|
|
addBody, _ := json.Marshal(addReq)
|
|
|
|
addHTTP, _ := http.NewRequest("POST", setup.Server.URL+"/api/collections/"+collectionID+"/books", bytes.NewBuffer(addBody))
|
|
addHTTP.Header.Set("Content-Type", "application/json")
|
|
addHTTP.Header.Set("Authorization", "Bearer "+setup.Token)
|
|
|
|
addResp, err := client.Do(addHTTP)
|
|
require.NoError(t, err)
|
|
defer addResp.Body.Close()
|
|
require.Equal(t, http.StatusNoContent, addResp.StatusCode)
|
|
|
|
// Admin should receive collection_updated message
|
|
msgAdmin := readWebSocketMessage(t, wsAdmin, 2*time.Second)
|
|
if msgAdmin["type"] != "collection_updated" {
|
|
t.Errorf("Admin should receive collection_updated, got %s", msgAdmin["type"])
|
|
}
|
|
|
|
// Regular user should NOT receive collection_updated message
|
|
wsRegular.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
|
|
_, _, err = wsRegular.ReadMessage()
|
|
if err == nil {
|
|
t.Errorf("Regular user should not receive collection_updated message")
|
|
}
|
|
}
|
|
|
|
// Helper: connectWebSocketToServer establishes WebSocket connection with auth token
|
|
func connectWebSocketToServer(t *testing.T, serverURL string, token string) *websocket.Conn {
|
|
wsURL := "ws" + strings.TrimPrefix(serverURL, "http") + "/ws/sync?token=" + token
|
|
|
|
ws, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
require.NoError(t, err, "WebSocket connection should succeed")
|
|
if resp != nil {
|
|
resp.Body.Close()
|
|
}
|
|
require.NotNil(t, ws, "WebSocket connection should be established")
|
|
|
|
// Wait for connection to be ready
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
return ws
|
|
}
|
|
|
|
// Helper: readWebSocketMessage reads a message from WebSocket with timeout
|
|
func readWebSocketMessage(t *testing.T, ws *websocket.Conn, timeout time.Duration) map[string]interface{} {
|
|
ws.SetReadDeadline(time.Now().Add(timeout))
|
|
|
|
_, message, err := ws.ReadMessage()
|
|
require.NoError(t, err, "Should receive WebSocket message")
|
|
|
|
var msg map[string]interface{}
|
|
err = json.Unmarshal(message, &msg)
|
|
require.NoError(t, err, "Message should be valid JSON")
|
|
|
|
return msg
|
|
}
|