diff --git a/cmd/server/tests/device_cap_test.go b/cmd/server/tests/device_cap_test.go
index 92a846f..5cff6cf 100644
--- a/cmd/server/tests/device_cap_test.go
+++ b/cmd/server/tests/device_cap_test.go
@@ -65,7 +65,7 @@ func TestUpdateUserMaxDevices(t *testing.T) {
 			}
 			jsonData, _ := json.Marshal(payload)
 
-			req, _ := http.NewRequest("PUT", ts.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
+			req, _ := http.NewRequest("PUT", setup.Server.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
 			req.Header.Set("Content-Type", "application/json")
 			req.Header.Set("Authorization", "Bearer "+adminToken)
 
@@ -131,7 +131,7 @@ func TestUpdateUserMaxDevicesValidation(t *testing.T) {
 			}
 			jsonData, _ := json.Marshal(payload)
 
-			req, _ := http.NewRequest("PUT", ts.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
+			req, _ := http.NewRequest("PUT", setup.Server.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
 			req.Header.Set("Content-Type", "application/json")
 			req.Header.Set("Authorization", "Bearer "+adminToken)
 
@@ -164,7 +164,7 @@ func TestUpdateUserMaxDevicesAuth(t *testing.T) {
 		}
 		jsonData, _ := json.Marshal(payload)
 
-		req, _ := http.NewRequest("PUT", ts.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
+		req, _ := http.NewRequest("PUT", setup.Server.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
 		req.Header.Set("Content-Type", "application/json")
 
 		client := &http.Client{}
@@ -187,7 +187,7 @@ func TestUpdateUserMaxDevicesAuth(t *testing.T) {
 		}
 		jsonData, _ := json.Marshal(payload)
 
-		req, _ := http.NewRequest("PUT", ts.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
+		req, _ := http.NewRequest("PUT", setup.Server.URL+"/api/auth/users/"+userID.String()+"/max-devices", bytes.NewBuffer(jsonData))
 		req.Header.Set("Content-Type", "application/json")
 		req.Header.Set("Authorization", "Bearer "+regularToken)
 
@@ -218,7 +218,7 @@ func TestUpdateUserMaxDevicesNonExistentUser(t *testing.T) {
 	}
 	jsonData, _ := json.Marshal(payload)
 
-	req, _ := http.NewRequest("PUT", ts.URL+"/api/auth/users/"+nonExistentUserID.String()+"/max-devices", bytes.NewBuffer(jsonData))
+	req, _ := http.NewRequest("PUT", setup.Server.URL+"/api/auth/users/"+nonExistentUserID.String()+"/max-devices", bytes.NewBuffer(jsonData))
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("Authorization", "Bearer "+adminToken)
 
@@ -247,7 +247,7 @@ func TestUpdateUserMaxDevicesMissingUserID(t *testing.T) {
 	jsonData, _ := json.Marshal(payload)
 
 	// Missing user ID in URL
-	req, _ := http.NewRequest("PUT", ts.URL+"/api/auth/users//max-devices", bytes.NewBuffer(jsonData))
+	req, _ := http.NewRequest("PUT", setup.Server.URL+"/api/auth/users//max-devices", bytes.NewBuffer(jsonData))
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("Authorization", "Bearer "+adminToken)
 
@@ -269,7 +269,7 @@ func TestListUsersIncludesMaxDevices(t *testing.T) {
 	adminUserID := getTestUserID(t, setup.DB)
 	adminToken = getAdminToken(t, setup.Server, adminUserID)
 
-	req, _ := http.NewRequest("GET", ts.URL+"/api/auth/users", nil)
+	req, _ := http.NewRequest("GET", setup.Server.URL+"/api/auth/users", nil)
 	req.Header.Set("Authorization", "Bearer "+adminToken)
 
 	client := &http.Client{}
diff --git a/cmd/server/tests/queue_test.go b/cmd/server/tests/queue_test.go
index c15354b..b53d947 100644
--- a/cmd/server/tests/queue_test.go
+++ b/cmd/server/tests/queue_test.go
@@ -41,7 +41,7 @@ func TestGetDeviceQueueStats(t *testing.T) {
 	userID := getTestUserID(t, setup.DB)
 
 	deviceToken := fmt.Sprintf("dev_%s", uuid.New().String())
-	device, err := db.CreateDevice(context.Background(), database.CreateDeviceParams{
+	device, err := setup.DB.CreateDevice(context.Background(), database.CreateDeviceParams{
 		UserID:           pgtype.UUID{Bytes: [16]byte(userID), Valid: true},
 		DeviceName:       "Test Device",
 		DeviceType:       "koreader",
@@ -77,7 +77,7 @@ func TestListDeviceQueueItems(t *testing.T) {
 	userID := getTestUserID(t, setup.DB)
 
 	deviceToken := fmt.Sprintf("dev_%s", uuid.New().String())
-	device, err := db.CreateDevice(context.Background(), database.CreateDeviceParams{
+	device, err := setup.DB.CreateDevice(context.Background(), database.CreateDeviceParams{
 		UserID:           pgtype.UUID{Bytes: [16]byte(userID), Valid: true},
 		DeviceName:       "Test Device",
 		DeviceType:       "koreader",
@@ -143,7 +143,7 @@ func TestClearDeviceQueue(t *testing.T) {
 	userID := getTestUserID(t, setup.DB)
 
 	deviceToken := fmt.Sprintf("dev_%s", uuid.New().String())
-	device, err := db.CreateDevice(context.Background(), database.CreateDeviceParams{
+	device, err := setup.DB.CreateDevice(context.Background(), database.CreateDeviceParams{
 		UserID:           pgtype.UUID{Bytes: [16]byte(userID), Valid: true},
 		DeviceName:       "Test Device",
 		DeviceType:       "koreader",
@@ -223,7 +223,7 @@ func loginAdminUser(t *testing.T, ts *httptest.Server, db *database.Queries) str
 
 	_, err = db.ListUsers(context.Background())
 	if err == nil {
-		return loginTestUser(t, setup.Server, setup.DB)
+		return loginTestUser(t, ts, db)
 	}
 
 	loginRequest := map[string]interface{}{
@@ -232,7 +232,7 @@ func loginAdminUser(t *testing.T, ts *httptest.Server, db *database.Queries) str
 	}
 	body, _ := json.Marshal(loginRequest)
 
-	req, _ := http.NewRequest("POST", setup.Server.URL+"/api/auth/login", bytes.NewBuffer(body))
+	req, _ := http.NewRequest("POST", ts.URL+"/api/auth/login", bytes.NewBuffer(body))
 	req.Header.Set("Content-Type", "application/json")
 
 	client := &http.Client{}
@@ -255,7 +255,7 @@ func loginUserWithID(t *testing.T, ts *httptest.Server, db *database.Queries, us
 	}
 	body, _ := json.Marshal(loginRequest)
 
-	req, _ := http.NewRequest("POST", setup.Server.URL+"/api/auth/login", bytes.NewBuffer(body))
+	req, _ := http.NewRequest("POST", ts.URL+"/api/auth/login", bytes.NewBuffer(body))
 	req.Header.Set("Content-Type", "application/json")
 
 	client := &http.Client{}
diff --git a/cmd/server/tests/websocket_test.go b/cmd/server/tests/websocket_test.go
index bb0bc9a..9861912 100644
--- a/cmd/server/tests/websocket_test.go
+++ b/cmd/server/tests/websocket_test.go
@@ -20,10 +20,10 @@ import (
 // TestWebSocketConnection tests basic WebSocket connection and authentication
 func TestWebSocketConnection(t *testing.T) {
 	// Setup test server with WebSocket
-	ts, queries, _ := setupTestServer(t)
+	setup := setupTestServer(t)
 
 	// Get JWT token for a test user
-	token := loginTestUser(t, ts, queries)
+	token := loginTestUser(t, setup.Server, setup.DB)
 
 	// Connect to WebSocket endpoint
 	wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
@@ -50,13 +50,13 @@ func TestWebSocketConnection(t *testing.T) {
 
 // TestWebSocketDeviceAuth tests device authentication via WebSocket
 func TestWebSocketDeviceAuth(t *testing.T) {
-	ts, queries, _ := setupTestServer(t)
+	setup := setupTestServer(t)
 
 	// Create a test device
-	userID := getTestUserID(t, queries)
+	userID := getTestUserID(t, setup.DB)
 	deviceID := uuid.New()
 
-	_, err := queries.CreateDevice(context.Background(), database.CreateDeviceParams{
+	_, err := setup.DB.CreateDevice(context.Background(), database.CreateDeviceParams{
 		UserID:               pgtype.UUID{Bytes: [16]byte(userID), Valid: true},
 		DeviceName:           "Test KOReader",
 		DeviceType:           "koreader",
@@ -76,21 +76,21 @@ func TestWebSocketDeviceAuth(t *testing.T) {
 
 	// We can't easily test WebSocket with custom headers using gorilla/websocket
 	// So this test just verifies the device exists
-	device, err := queries.GetDeviceByAuthToken(context.Background(), "test-device-token-"+deviceID.String())
+	device, err := setup.DB.GetDeviceByAuthToken(context.Background(), "test-device-token-"+deviceID.String())
 	require.NoError(t, err)
 	assert.Equal(t, "Test KOReader", device.DeviceName)
 }
 
 // TestWebSocketProgressBroadcast tests that progress updates are broadcast to connected clients
 func TestWebSocketProgressBroadcast(t *testing.T) {
-	ts, queries, _ := setupTestServer(t)
+	setup := setupTestServer(t)
 
 	// Get JWT token
-	token := loginTestUser(t, ts, queries)
+	token := loginTestUser(t, setup.Server, setup.DB)
 
 	// Create a test media item
-	userID := getTestUserID(t, queries)
-	mediaID := createTestMediaItem(t, queries, userID)
+	userID := getTestUserID(t, setup.DB)
+	mediaID := createTestMediaItem(t, setup.DB, userID)
 
 	// Connect WebSocket client
 	wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
@@ -145,9 +145,9 @@ func TestWebSocketProgressBroadcast(t *testing.T) {
 
 // TestWebSocketPingPong tests that ping/pong messages work correctly
 func TestWebSocketPingPong(t *testing.T) {
-	ts, queries, _ := setupTestServer(t)
+	setup := setupTestServer(t)
 
-	token := loginTestUser(t, ts, queries)
+	token := loginTestUser(t, setup.Server, setup.DB)
 
 	wsURL := strings.Replace(setup.Server.URL, "http", "ws", 1) + "/ws/sync?token=" + token
 	ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
@@ -177,9 +177,9 @@ func TestWebSocketPingPong(t *testing.T) {
 
 // TestWebSocketConnectionLimit tests that the server handles multiple connections
 func TestWebSocketConnectionLimit(t *testing.T) {
-	ts, queries, _ := setupTestServer(t)
+	setup := setupTestServer(t)
 
-	token := loginTestUser(t, ts, queries)
+	token := loginTestUser(t, setup.Server, setup.DB)
 
 	// Create multiple connections
 	connections := make([]*websocket.Conn, 5)