package httpapi import ( "context" "net/http" "net/http/httptest" "testing" "time" "github.com/coder/websocket" "github.com/vasyansk/imap-copier/internal/config" "github.com/vasyansk/imap-copier/internal/crypto" "github.com/vasyansk/imap-copier/internal/wshub" ) func TestWSRequiresAuth(t *testing.T) { s := &Server{cfg: config.Config{SessionSecret: []byte("x")}, hub: wshub.New()} srv := httptest.NewServer(s.Router()) defer srv.Close() // no cookie -> upgrade rejected (401) _, resp, err := websocket.Dial(context.Background(), "ws"+srv.URL[4:]+"/ws?task_id=1", nil) if err == nil { t.Fatal("expected auth rejection") } if resp != nil && resp.StatusCode != http.StatusUnauthorized { t.Fatalf("want 401, got %d", resp.StatusCode) } } func TestWSUnsubscribesOnClientDisconnect(t *testing.T) { hub := wshub.New() secret := []byte("sekret") s := &Server{cfg: config.Config{AuthUser: "admin", SessionSecret: secret}, hub: hub} srv := httptest.NewServer(s.Router()) defer srv.Close() tok := crypto.SignSession(secret, "admin", time.Now().Add(time.Hour)) hdr := http.Header{} hdr.Set("Cookie", cookieName+"="+tok) ctx := context.Background() c, _, err := websocket.Dial(ctx, "ws"+srv.URL[4:]+"/ws?task_id=7", &websocket.DialOptions{HTTPHeader: hdr}) if err != nil { t.Fatalf("dial: %v", err) } // wait until subscribed deadline := time.Now().Add(2 * time.Second) for hub.SubscriberCount(7) == 0 { if time.Now().After(deadline) { t.Fatal("never subscribed") } time.Sleep(10 * time.Millisecond) } // abrupt client close -> server must detect and unsubscribe c.CloseNow() deadline = time.Now().Add(3 * time.Second) for hub.SubscriberCount(7) != 0 { if time.Now().After(deadline) { t.Fatal("subscription leaked after client disconnect") } time.Sleep(20 * time.Millisecond) } }