summaryrefslogtreecommitdiff
path: root/internal/api/websocket_test.go
diff options
context:
space:
mode:
authorPeter Stone <thepeterstone@gmail.com>2026-03-08 21:03:50 +0000
committerPeter Stone <thepeterstone@gmail.com>2026-03-08 21:03:50 +0000
commit632ea5a44731af94b6238f330a3b5440906c8ae7 (patch)
treed8c780412598d66b89ef390b5729e379fdfd9d5b /internal/api/websocket_test.go
parent406247b14985ab57902e8e42898dc8cb8960290d (diff)
parent93a4c852bf726b00e8014d385165f847763fa214 (diff)
merge: pull latest from master and resolve conflicts
- Resolve conflicts in API server, CLI, and executor. - Maintain Gemini classification and assignment logic. - Update UI to use generic agent config and project_dir. - Fix ProjectDir/WorkingDir inconsistencies in Gemini runner. - All tests passing after merge.
Diffstat (limited to 'internal/api/websocket_test.go')
-rw-r--r--internal/api/websocket_test.go221
1 files changed, 221 insertions, 0 deletions
diff --git a/internal/api/websocket_test.go b/internal/api/websocket_test.go
new file mode 100644
index 0000000..72b83f2
--- /dev/null
+++ b/internal/api/websocket_test.go
@@ -0,0 +1,221 @@
+package api
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+ "time"
+
+ "golang.org/x/net/websocket"
+)
+
+// TestWebSocket_RejectsConnectionWithoutToken verifies that when an API token
+// is configured, WebSocket connections without a valid token are rejected with 401.
+func TestWebSocket_RejectsConnectionWithoutToken(t *testing.T) {
+ srv, _ := testServer(t)
+ srv.SetAPIToken("secret-token")
+
+ // Plain HTTP request simulates a WebSocket upgrade attempt without token.
+ req := httptest.NewRequest("GET", "/api/ws", nil)
+ w := httptest.NewRecorder()
+ srv.Handler().ServeHTTP(w, req)
+
+ if w.Code != http.StatusUnauthorized {
+ t.Errorf("want 401, got %d", w.Code)
+ }
+}
+
+// TestWebSocket_RejectsConnectionWithWrongToken verifies a wrong token is rejected.
+func TestWebSocket_RejectsConnectionWithWrongToken(t *testing.T) {
+ srv, _ := testServer(t)
+ srv.SetAPIToken("secret-token")
+
+ req := httptest.NewRequest("GET", "/api/ws?token=wrong-token", nil)
+ w := httptest.NewRecorder()
+ srv.Handler().ServeHTTP(w, req)
+
+ if w.Code != http.StatusUnauthorized {
+ t.Errorf("want 401, got %d", w.Code)
+ }
+}
+
+// TestWebSocket_AcceptsConnectionWithValidQueryToken verifies a valid token in
+// the query string is accepted.
+func TestWebSocket_AcceptsConnectionWithValidQueryToken(t *testing.T) {
+ srv, _ := testServer(t)
+ srv.SetAPIToken("secret-token")
+ srv.StartHub()
+ ts := httptest.NewServer(srv.Handler())
+ t.Cleanup(ts.Close)
+
+ wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/api/ws?token=secret-token"
+ ws, err := websocket.Dial(wsURL, "", "http://localhost/")
+ if err != nil {
+ t.Fatalf("expected connection to succeed with valid token: %v", err)
+ }
+ ws.Close()
+}
+
+// TestWebSocket_AcceptsConnectionWithBearerToken verifies a valid token in the
+// Authorization header is accepted.
+func TestWebSocket_AcceptsConnectionWithBearerToken(t *testing.T) {
+ srv, _ := testServer(t)
+ srv.SetAPIToken("secret-token")
+ srv.StartHub()
+ ts := httptest.NewServer(srv.Handler())
+ t.Cleanup(ts.Close)
+
+ wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/api/ws"
+ cfg, err := websocket.NewConfig(wsURL, "http://localhost/")
+ if err != nil {
+ t.Fatalf("config: %v", err)
+ }
+ cfg.Header = http.Header{"Authorization": {"Bearer secret-token"}}
+ ws, err := websocket.DialConfig(cfg)
+ if err != nil {
+ t.Fatalf("expected connection to succeed with Bearer token: %v", err)
+ }
+ ws.Close()
+}
+
+// TestWebSocket_NoTokenConfigured verifies that when no API token is set,
+// connections are allowed without authentication.
+func TestWebSocket_NoTokenConfigured(t *testing.T) {
+ srv, _ := testServer(t)
+ // No SetAPIToken call — auth is disabled.
+ srv.StartHub()
+ ts := httptest.NewServer(srv.Handler())
+ t.Cleanup(ts.Close)
+
+ wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/api/ws"
+ ws, err := websocket.Dial(wsURL, "", "http://localhost/")
+ if err != nil {
+ t.Fatalf("expected connection without token when auth disabled: %v", err)
+ }
+ ws.Close()
+}
+
+// TestWebSocket_RejectsConnectionWhenAtMaxClients verifies that when the hub
+// is at capacity, new WebSocket upgrade requests are rejected with 503.
+func TestWebSocket_RejectsConnectionWhenAtMaxClients(t *testing.T) {
+ orig := maxWsClients
+ maxWsClients = 0 // immediately at capacity
+ t.Cleanup(func() { maxWsClients = orig })
+
+ srv, _ := testServer(t)
+ srv.StartHub()
+
+ req := httptest.NewRequest("GET", "/api/ws", nil)
+ w := httptest.NewRecorder()
+ srv.Handler().ServeHTTP(w, req)
+
+ if w.Code != http.StatusServiceUnavailable {
+ t.Errorf("want 503, got %d", w.Code)
+ }
+}
+
+// TestWebSocket_StaleConnectionCleanedUp verifies that when a client
+// disconnects (or the connection is closed), the hub unregisters it.
+// Short ping intervals are used so the test completes quickly.
+func TestWebSocket_StaleConnectionCleanedUp(t *testing.T) {
+ origInterval := wsPingInterval
+ origDeadline := wsPingDeadline
+ wsPingInterval = 20 * time.Millisecond
+ wsPingDeadline = 20 * time.Millisecond
+ t.Cleanup(func() {
+ wsPingInterval = origInterval
+ wsPingDeadline = origDeadline
+ })
+
+ srv, _ := testServer(t)
+ srv.StartHub()
+ ts := httptest.NewServer(srv.Handler())
+ t.Cleanup(ts.Close)
+
+ wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/api/ws"
+ ws, err := websocket.Dial(wsURL, "", "http://localhost/")
+ if err != nil {
+ t.Fatalf("dial: %v", err)
+ }
+
+ // Wait for hub to register the client.
+ deadline := time.Now().Add(200 * time.Millisecond)
+ for time.Now().Before(deadline) {
+ if srv.hub.ClientCount() == 1 {
+ break
+ }
+ time.Sleep(5 * time.Millisecond)
+ }
+ if got := srv.hub.ClientCount(); got != 1 {
+ t.Fatalf("before close: want 1 client, got %d", got)
+ }
+
+ // Close connection without a proper WebSocket close handshake
+ // to simulate a client crash / network drop.
+ ws.Close()
+
+ // Hub should unregister the client promptly.
+ deadline = time.Now().Add(500 * time.Millisecond)
+ for time.Now().Before(deadline) {
+ if srv.hub.ClientCount() == 0 {
+ return
+ }
+ time.Sleep(10 * time.Millisecond)
+ }
+ t.Fatalf("after close: expected 0 clients, got %d", srv.hub.ClientCount())
+}
+
+// TestWebSocket_PingWriteDeadlineEvictsStaleConn verifies that a stale
+// connection (write times out) is eventually evicted by the ping goroutine.
+// It uses a very short write deadline to force a timeout on a connection
+// whose receive buffer is full.
+func TestWebSocket_PingWriteDeadlineEvictsStaleConn(t *testing.T) {
+ origInterval := wsPingInterval
+ origDeadline := wsPingDeadline
+ // Very short deadline: ping fails almost immediately after the first tick.
+ wsPingInterval = 30 * time.Millisecond
+ wsPingDeadline = 1 * time.Millisecond
+ t.Cleanup(func() {
+ wsPingInterval = origInterval
+ wsPingDeadline = origDeadline
+ })
+
+ srv, _ := testServer(t)
+ srv.StartHub()
+ ts := httptest.NewServer(srv.Handler())
+ t.Cleanup(ts.Close)
+
+ wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/api/ws"
+ ws, err := websocket.Dial(wsURL, "", "http://localhost/")
+ if err != nil {
+ t.Fatalf("dial: %v", err)
+ }
+ defer ws.Close()
+
+ // Wait for registration.
+ deadline := time.Now().Add(200 * time.Millisecond)
+ for time.Now().Before(deadline) {
+ if srv.hub.ClientCount() == 1 {
+ break
+ }
+ time.Sleep(5 * time.Millisecond)
+ }
+ if got := srv.hub.ClientCount(); got != 1 {
+ t.Fatalf("before stale: want 1 client, got %d", got)
+ }
+
+ // The connection itself is alive (loopback), so the 1ms deadline is generous
+ // enough to succeed. This test mainly verifies the ping goroutine doesn't
+ // panic and that ClientCount stays consistent after disconnect.
+ ws.Close()
+
+ deadline = time.Now().Add(500 * time.Millisecond)
+ for time.Now().Before(deadline) {
+ if srv.hub.ClientCount() == 0 {
+ return
+ }
+ time.Sleep(10 * time.Millisecond)
+ }
+ t.Fatalf("expected 0 clients after stale eviction, got %d", srv.hub.ClientCount())
+}