package wsclient import ( "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/gorilla/websocket" ) // startWSServer brings up a minimal WS test server. The server // reads the auth frame, replies with `{"ready": true}`, then // echoes every subsequent frame back as a "ready" ack. We use // this to assert Connect's frame-level behavior without // pulling in the full ingestd path. func startWSServer(t *testing.T, accept func(auth []byte) (reply []byte, ok bool)) (string, *httptest.Server) { t.Helper() upgrader := websocket.Upgrader{} srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { c, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Log("upgrade:", err) return } defer c.Close() // auth frame _ = c.SetReadDeadline(time.Now().Add(2 * time.Second)) _, msg, err := c.ReadMessage() if err != nil { return } reply, ok := accept(msg) if !ok { return } _ = c.SetWriteDeadline(time.Now().Add(2 * time.Second)) _ = c.WriteMessage(websocket.TextMessage, reply) // echo loop for { _ = c.SetReadDeadline(time.Now().Add(2 * time.Second)) _, body, err := c.ReadMessage() if err != nil { return } _ = c.SetWriteDeadline(time.Now().Add(2 * time.Second)) _ = c.WriteMessage(websocket.TextMessage, []byte(`{"echo":`+string(body)+`}`)) } })) wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" return wsURL, srv } func TestConnect_RoundTrip(t *testing.T) { wsURL, srv := startWSServer(t, func(auth []byte) ([]byte, bool) { // Expect {"api_key":"..."} and reply with ready. if !strings.Contains(string(auth), `"api_key"`) { return []byte(`{"error":"bad auth"}`), false } return []byte(`{"ready":true}`), true }) defer srv.Close() c, err := Connect(Config{URL: wsURL, APIKey: "acme-001:prom-prod:s3cret-acme"}) if err != nil { t.Fatal("connect:", err) } defer c.Close() if got := string(c.AuthReply()); got != `{"ready":true}` { t.Fatalf("auth reply=%q, want ready", got) } ack, err := c.SendAlert([]byte(`{"hello":"world"}`)) if err != nil { t.Fatal("send:", err) } if !strings.Contains(string(ack), `"echo":{"hello":"world"}`) { t.Fatalf("echo ack=%q, want echo", ack) } } func TestConnect_BadAuth(t *testing.T) { // Server accepts the dial, reads the auth frame, then // closes. Connect should see the close on the read and // return the error. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { c, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) if err != nil { return } _ = c.SetReadDeadline(time.Now().Add(2 * time.Second)) _, _, _ = c.ReadMessage() _ = c.WriteMessage(websocket.TextMessage, []byte(`{"error":"unauthorized"}`)) _ = c.Close() })) defer srv.Close() wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" c, err := Connect(Config{URL: wsURL, APIKey: "x"}) if err != nil { t.Fatal("connect (first frame should still round-trip):", err) } defer c.Close() if !strings.Contains(string(c.AuthReply()), `"error"`) { t.Fatalf("expected error in auth reply, got %q", c.AuthReply()) } } func TestConnect_RejectsEmptyConfig(t *testing.T) { if _, err := Connect(Config{URL: "", APIKey: "x"}); err == nil { t.Fatal("empty URL should error") } if _, err := Connect(Config{URL: "ws://x", APIKey: ""}); err == nil { t.Fatal("empty APIKey should error") } }