diff --git a/cmd/server/tests/test_helpers.go b/cmd/server/tests/test_helpers.go index 3c25b83..6b0e85b 100644 --- a/cmd/server/tests/test_helpers.go +++ b/cmd/server/tests/test_helpers.go @@ -8,7 +8,7 @@ import ( ratelimit "bookhoard/internal/middleware" "bookhoard/internal/router" "bookhoard/internal/services" - "bookhoard/internal/sync" + wsync "bookhoard/internal/sync" "bytes" "context" "encoding/json" @@ -17,6 +17,7 @@ import ( "net/http/httptest" "os" "strings" + "sync" "testing" "time" @@ -65,6 +66,58 @@ type DeviceTestData struct { PGType database.Devices } +// TestServerSetup manages the lifecycle of a test server with proper resource cleanup +type TestServerSetup struct { + Server *httptest.Server + DB *database.Queries + DBPool *pgxpool.Pool + Config *config.Config + ConnManager *wsync.ConnectionManager + QueueProcessor *wsync.SyncQueueProcessor + CleanupCancel context.CancelFunc + QueueCtx context.Context + QueueCancel context.CancelFunc + mu sync.Mutex + closed bool +} + +// Close cleans up all resources in the correct order +func (s *TestServerSetup) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return nil + } + + // Stop queue processor first + if s.QueueCancel != nil { + s.QueueCancel() + s.QueueCancel = nil + } + + // Stop connection manager cleanup task + if s.CleanupCancel != nil { + s.CleanupCancel() + s.CleanupCancel = nil + } + + // Close HTTP server + if s.Server != nil { + s.Server.Close() + s.Server = nil + } + + // Close database pool (this waits for all connections to be released) + if s.DBPool != nil { + s.DBPool.Close() + s.DBPool = nil + } + + s.closed = true + return nil +} + // Helper functions for testing func containsPrefix(s, prefix string) bool { return len(s) >= len(prefix) && s[:len(prefix)] == prefix @@ -130,18 +183,18 @@ func getCachePath() string { // setupDeviceTest creates a complete test environment for device tests func setupDeviceTest(t *testing.T) *TestDeviceSetup { - ts, db, cfg := setupTestServer(t) + serverSetup := setupTestServer(t) // Create user ONCE with known credentials - user := createTestUserOnce(t, db) + user := createTestUserOnce(t, serverSetup.DB) // Login to get token - token := loginUserWithCredentials(t, ts, user.Email, user.Password) + token := loginUserWithCredentials(t, serverSetup.Server, user.Email, user.Password) return &TestDeviceSetup{ - Server: ts, - DB: db, - Config: cfg, + Server: serverSetup.Server, + DB: serverSetup.DB, + Config: serverSetup.Config, User: user, UserToken: token, } @@ -241,8 +294,8 @@ func (s *TestDeviceSetup) CreateDevice(t *testing.T, deviceName, deviceType, dev } // setupTestServer creates a test server with a test database -// Returns: (*httptest.Server, *database.Queries, *config.Config) -func setupTestServer(t *testing.T) (*httptest.Server, *database.Queries, *config.Config) { +// Returns: *TestServerSetup with automatic cleanup via t.Cleanup +func setupTestServer(t *testing.T) *TestServerSetup { // Load configuration using the same method as main application cfg := config.LoadConfig() @@ -271,12 +324,13 @@ func setupTestServer(t *testing.T) (*httptest.Server, *database.Queries, *config deviceAuthMiddleware := middleware.NewDeviceAuthMiddleware(queries) // Create WebSocket connection manager - connManager := sync.NewConnectionManager() - connManager.StartCleanupTask() + connManager := wsync.NewConnectionManager() + cleanupCancel := connManager.StartCleanupTask() - // Create sync queue processor - queueProcessor := sync.NewSyncQueueProcessor(queries) - go queueProcessor.Start(context.Background()) + // Create sync queue processor with cancellable context + queueProcessor := wsync.NewSyncQueueProcessor(queries) + queueCtx, queueCancel := context.WithCancel(context.Background()) + go queueProcessor.Start(queueCtx) koreaderHandler := handlers.NewKOReaderHandler(queries, connManager, queueProcessor) wsHandler := handlers.NewWSHandler(queries, connManager, cfg.JWTSecret, deviceAuthMiddleware) @@ -341,8 +395,27 @@ func setupTestServer(t *testing.T) (*httptest.Server, *database.Queries, *config // Create test server ts := httptest.NewServer(e) - // Return server, queries, and config - return ts, queries, cfg + // Create TestServerSetup struct with all resources + setup := &TestServerSetup{ + Server: ts, + DB: queries, + DBPool: dbPool, + Config: cfg, + ConnManager: connManager, + QueueProcessor: queueProcessor, + CleanupCancel: cleanupCancel, + QueueCtx: queueCtx, + QueueCancel: queueCancel, + } + + // Register cleanup function to run automatically when test completes + t.Cleanup(func() { + if err := setup.Close(); err != nil { + t.Errorf("Failed to cleanup test server: %v", err) + } + }) + + return setup } // loginTestUser logs in a test user and returns the JWT token