Something went wrong. Try again.
Monorepo for Tangled forked from tangled.org/core
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322package 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") }}