package jetstream
import (
"context"
"errors"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"sync"
"testing"
"time"
"github.com/gorilla/websocket"
"github.com/klauspost/compress/zstd"
)
type streamServer struct {
*httptest.Server
mu sync.Mutex
queries []url.Values
headers []http.Header
}
func newStreamServer(t *testing.T, script func(n int, con *websocket.Conn)) *streamServer {
t.Helper()
s := &streamServer{}
s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
con, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
if err != nil {
return
}
defer con.Close()
s.mu.Lock()
s.queries = append(s.queries, r.URL.Query())
s.headers = append(s.headers, r.Header)
n := len(s.queries) - 1
s.mu.Unlock()
script(n, con)
}))
t.Cleanup(s.Close)
return s
}
func (s *streamServer) endpoint() string {
return "ws" + strings.TrimPrefix(s.URL, "http") + "/subscribe"
}
func (s *streamServer) query(n int) url.Values {
s.mu.Lock()
defer s.mu.Unlock()
return s.queries[n]
}
func (s *streamServer) connects() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.queries)
}
func drain(con *websocket.Conn) {
for {
if _, _, err := con.ReadMessage(); err != nil {
return
}
}
}
func testStream(t *testing.T, cfg StreamConfig) *Stream {
t.Helper()
if cfg.Cursor == nil {
cfg.Cursor = func() (int64, bool, error) { return 0, false, nil }
}
if cfg.ReconnectDelay == 0 {
cfg.ReconnectDelay = 10 * time.Millisecond
}
s, err := NewStream(cfg, slog.New(slog.NewTextHandler(io.Discard, nil)))
if err != nil {
t.Fatal(err)
}
return s
}
func run(t *testing.T, s *Stream, until func(*Event) bool) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
err := make(chan error, 1)
go func() {
err <- s.Run(ctx, func(_ context.Context, e *Event) error {
if until(e) {
cancel()
}
return nil
})
}()
if e := <-err; e != nil {
t.Fatal(e)
}
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
t.Fatal("the stream never delivered what the test waited for")
}
}
func TestStreamURLCarriesEachParamOnce(t *testing.T) {
s := testStream(t, StreamConfig{
Endpoint: "ws://hydrant/subscribe?cursor=9&wantedEventTypes=old",
Collections: []string{"org.tangled.track.ticket", "org.tangled.feed.comment"},
EventTypes: []string{"live", "historical"},
})
got, err := s.url(new(int64))
if err != nil {
t.Fatal(err)
}
u, err := url.Parse(got)
if err != nil {
t.Fatal(err)
}
q := u.Query()
if strings.Count(got, "cursor=") != 1 || q.Get("cursor") != "0" {
t.Errorf("url %q should carry exactly one cursor=0", got)
}
if strings.Count(got, "?") != 1 {
t.Errorf("url %q has more than one query string", got)
}
if want := []string{"live", "historical"}; strings.Join(q["wantedEventTypes"], ",") != strings.Join(want, ",") {
t.Errorf("wantedEventTypes = %v, want %v", q["wantedEventTypes"], want)
}
if len(q["wantedCollections"]) != 2 {
t.Errorf("wantedCollections = %v, want both", q["wantedCollections"])
}
}
func TestStreamDecodesLive(t *testing.T) {
srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
for _, frame := range []string{
`{"did":"d","time_us":1,"kind":"commit","commit":{"rev":"r","operation":"create","live":true}}`,
`{"did":"d","time_us":2,"kind":"commit","commit":{"rev":"r","operation":"create","live":false}}`,
`{"did":"d","time_us":3,"kind":"commit","commit":{"rev":"r","operation":"create"}}`,
} {
con.WriteMessage(websocket.TextMessage, []byte(frame))
}
drain(con)
})
s := testStream(t, StreamConfig{Endpoint: srv.endpoint()})
live := map[int64]bool{}
run(t, s, func(e *Event) bool {
live[e.TimeUS] = e.Commit.IsLive()
return e.TimeUS == 3
})
if want := map[int64]bool{1: true, 2: false, 3: true}; len(live) != 3 || live[1] != want[1] || live[2] != want[2] || live[3] != want[3] {
t.Fatalf("live = %v, want %v (absent means live)", live, want)
}
}
func TestStreamSkipsUndecodableFrames(t *testing.T) {
srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
con.WriteMessage(websocket.TextMessage, []byte(`not json`))
con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":7}`))
drain(con)
})
run(t, testStream(t, StreamConfig{Endpoint: srv.endpoint()}), func(e *Event) bool { return e.TimeUS == 7 })
}
func TestStreamResumesFromTheHandlersCursorRewound(t *testing.T) {
var mu sync.Mutex
cursor, have := int64(0), false
srv := newStreamServer(t, func(n int, con *websocket.Conn) {
if n == 0 {
con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":10000000}`))
}
if n == 1 {
con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":10000001}`))
}
drain(con)
})
s := testStream(t, StreamConfig{
Endpoint: srv.endpoint(),
Rewind: 2 * time.Second,
Cursor: func() (int64, bool, error) {
mu.Lock()
defer mu.Unlock()
return cursor, have, nil
},
})
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
done := make(chan struct{})
go func() {
defer close(done)
s.Run(ctx, func(_ context.Context, e *Event) error {
mu.Lock()
cursor, have = e.TimeUS, true
mu.Unlock()
if e.TimeUS == 10_000_000 {
return errors.New("drop the connection")
}
cancel()
return nil
})
}()
<-done
if got := srv.query(0).Get("cursor"); got != "0" {
t.Errorf("empty db cursor = %q, want 0 so a fresh db replays", got)
}
if got := srv.query(1).Get("cursor"); got != "8000000" {
t.Errorf("reconnect cursor = %q, want the db cursor minus the rewind", got)
}
}
func TestStreamOnlyReadsTheCursorOncePerConnect(t *testing.T) {
var reads int
srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":5}`))
drain(con)
})
s := testStream(t, StreamConfig{
Endpoint: srv.endpoint(),
Cursor: func() (int64, bool, error) { reads++; return 0, false, nil },
})
run(t, s, func(e *Event) bool { return true })
if reads != 1 {
t.Fatalf("cursor source read %d times for one connection, want 1", reads)
}
}
func TestStreamCursorErrorStopsTheConnect(t *testing.T) {
srv := newStreamServer(t, func(int, *websocket.Conn) { t.Error("connected despite an unreadable cursor") })
s := testStream(t, StreamConfig{
Endpoint: srv.endpoint(),
Cursor: func() (int64, bool, error) { return 0, false, errors.New("db down") },
})
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
if err := s.Run(ctx, func(context.Context, *Event) error { return nil }); err != nil {
t.Fatal(err)
}
}
func TestStreamReadTimeoutReconnects(t *testing.T) {
srv := newStreamServer(t, func(n int, con *websocket.Conn) {
if n == 0 {
drain(con)
return
}
con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":9}`))
drain(con)
})
s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), ReadTimeout: 100 * time.Millisecond})
run(t, s, func(e *Event) bool { return e.TimeUS == 9 })
}
func TestStreamPingsKeepAConnectionAlive(t *testing.T) {
srv := newStreamServer(t, func(n int, con *websocket.Conn) {
for range 6 {
time.Sleep(60 * time.Millisecond)
con.WriteControl(websocket.PingMessage, nil, time.Now().Add(time.Second))
}
con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":4}`))
drain(con)
})
s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), ReadTimeout: 150 * time.Millisecond})
run(t, s, func(e *Event) bool { return e.TimeUS == 4 })
if n := len(srv.queries); n != 1 {
t.Fatalf("connected %d times, pings should have kept the first alive", n)
}
}
func TestStreamDecodesZstdWithTheJetstreamDictionary(t *testing.T) {
enc, err := zstd.NewWriter(nil, zstd.WithEncoderDict(zstdDictionary))
if err != nil {
t.Fatal(err)
}
srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
frame := enc.EncodeAll([]byte(`{"did":"d","time_us":6,"kind":"commit","commit":{"live":false}}`), nil)
con.WriteMessage(websocket.BinaryMessage, frame)
drain(con)
})
s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), Compress: true, UserAgent: "test"})
run(t, s, func(e *Event) bool { return e.TimeUS == 6 && !e.Commit.IsLive() })
if got := srv.headers[0].Get("Socket-Encoding"); got != "zstd" {
t.Errorf("Socket-Encoding = %q, want zstd", got)
}
if got := srv.headers[0].Get("User-Agent"); got != "test" {
t.Errorf("User-Agent = %q, want test", got)
}
}
func TestStreamDecodesZstdWithoutADictionary(t *testing.T) {
enc, err := zstd.NewWriter(nil)
if err != nil {
t.Fatal(err)
}
srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
frame := enc.EncodeAll([]byte(`{"did":"d","time_us":7,"kind":"commit","commit":{"live":false}}`), nil)
con.WriteMessage(websocket.BinaryMessage, frame)
drain(con)
})
s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), Compress: true})
run(t, s, func(e *Event) bool { return e.TimeUS == 7 && !e.Commit.IsLive() })
}
func TestStreamWithoutACursorTailsWhenAsked(t *testing.T) {
srv := newStreamServer(t, func(_ int, con *websocket.Conn) { drain(con) })
s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), TailWhenNoCursor: true})
ctx, cancel := context.WithCancel(context.Background())
go s.Run(ctx, func(context.Context, *Event) error { return nil })
for srv.connects() == 0 {
time.Sleep(time.Millisecond)
}
cancel()
if _, ok := srv.query(0)["cursor"]; ok {
t.Fatal("a stream with no cursor should connect without one")
}
}