broad-announce/internal/wsclient/client_test.go

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")
	}
}