1package agentrunner23import (4 "context"5 "encoding/json"6 "errors"7 "net/http"8 "net/http/httptest"9 "strings"10 "sync"11 "testing"12 "time"1314 "github.com/FelineStateMachine/tinyrelay/internal/event"15 "github.com/coder/websocket"16)1718const relayTestSecret = "0000000000000000000000000000000000000000000000000000000000000001"1920func TestResolveRootQueriesUnseenParentsWithoutDispatchingThem(t *testing.T) {21 room := "room"22 root := relayTestEvent(t, 11, room, "")23 middle := relayTestEvent(t, 12, room, root.ID)24 leaf := relayTestEvent(t, 12, room, middle.ID)25 events := map[string]event.Event{root.ID: root, middle.ID: middle, leaf.ID: leaf}26 var mu sync.Mutex27 var closed []string28 ready := make(chan struct{}, 1)29 server := relayTestServer(t, events, &mu, &closed, ready)30 defer server.Close()3132 client := RelayClient{URL: "ws" + strings.TrimPrefix(server.URL, "http"), Secret: relayTestSecret, Rooms: []string{room}}33 var dispatched int34 client.OnEvent = func(event.Event) { dispatched++ }35 ctx, cancel := context.WithCancel(context.Background())36 defer cancel()37 done := make(chan error, 1)38 go func() { done <- client.Run(ctx) }()39 select {40 case <-ready:41 case <-time.After(time.Second):42 t.Fatal("relay did not finish its initial subscription")43 }4445 queryCtx, queryCancel := context.WithTimeout(context.Background(), time.Second)46 defer queryCancel()47 got, err := client.ResolveRoot(queryCtx, room, leaf.ID)48 if err != nil {49 t.Fatal(err)50 }51 if got != root.ID {52 t.Fatalf("root=%s, want %s", got, root.ID)53 }54 if dispatched != 0 {55 t.Fatalf("query events dispatched to agent callback: %d", dispatched)56 }57 mu.Lock()58 closeCount := len(closed)59 mu.Unlock()60 deadline := time.Now().Add(time.Second)61 for closeCount < 3 && time.Now().Before(deadline) {62 time.Sleep(time.Millisecond)63 mu.Lock()64 closeCount = len(closed)65 mu.Unlock()66 }67 if closeCount != 3 {68 t.Fatalf("CLOSE frames=%d, want 3", closeCount)69 }70 cancel()71 select {72 case <-done:73 case <-time.After(time.Second):74 t.Fatal("relay did not stop")75 }76}7778func TestResolveRootRejectsInvalidQueriedEvents(t *testing.T) {79 tests := []struct {80 name string81 make func(t *testing.T) (string, event.Event)82 want string83 }{84 {"wrong room", func(t *testing.T) (string, event.Event) {85 e := relayTestEvent(t, 11, "other", "")86 return "room", e87 }, "another room"},88 {"invalid kind", func(t *testing.T) (string, event.Event) {89 e := relayTestEvent(t, 1, "room", "")90 return "room", e91 }, "invalid kind"},92 }93 for _, tt := range tests {94 t.Run(tt.name, func(t *testing.T) {95 room, e := tt.make(t)96 server := relayTestServer(t, map[string]event.Event{e.ID: e}, nil, nil, make(chan struct{}, 1))97 defer server.Close()98 client, stop := runningTestRelay(t, server, room)99 defer stop()100 ctx, cancel := context.WithTimeout(context.Background(), time.Second)101 defer cancel()102 _, err := client.ResolveRoot(ctx, room, e.ID)103 if err == nil || !strings.Contains(err.Error(), tt.want) {104 t.Fatalf("error=%v, want substring %q", err, tt.want)105 }106 })107 }108}109110func TestResolveRootTimesOutAndClosesQuery(t *testing.T) {111 root := relayTestEvent(t, 11, "room", "")112 server := relayTestServer(t, map[string]event.Event{}, nil, nil, make(chan struct{}, 1))113 defer server.Close()114 client, stop := runningTestRelay(t, server, "room")115 defer stop()116 ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)117 defer cancel()118 _, err := client.ResolveRoot(ctx, "room", root.ID)119 if !errors.Is(err, context.DeadlineExceeded) {120 t.Fatalf("error=%v, want deadline exceeded", err)121 }122}123124func relayTestEvent(t *testing.T, kind int, room, parent string) event.Event {125 t.Helper()126 tags := [][]string{{"h", room}}127 if parent != "" {128 tags = append(tags, []string{"e", parent, "", "root"})129 }130 e := event.Event{Kind: kind, CreatedAt: time.Now().Unix(), Tags: tags, Content: "test"}131 if err := event.Sign(&e, relayTestSecret); err != nil {132 t.Fatal(err)133 }134 return e135}136137func runningTestRelay(t *testing.T, server *httptest.Server, room string) (*RelayClient, func()) {138 t.Helper()139 client := &RelayClient{URL: "ws" + strings.TrimPrefix(server.URL, "http"), Secret: relayTestSecret, Rooms: []string{room}}140 ctx, cancel := context.WithCancel(context.Background())141 done := make(chan error, 1)142 go func() { done <- client.Run(ctx) }()143 deadline := time.Now().Add(time.Second)144 for {145 client.mu.Lock()146 connected := client.conn != nil147 client.mu.Unlock()148 if connected {149 break150 }151 if time.Now().After(deadline) {152 t.Fatal("relay did not connect")153 }154 time.Sleep(time.Millisecond)155 }156 return client, func() {157 cancel()158 select {159 case <-done:160 case <-time.After(time.Second):161 t.Fatal("relay did not stop")162 }163 }164}165166func relayTestServer(t *testing.T, events map[string]event.Event, mu *sync.Mutex, closed *[]string, ready chan struct{}) *httptest.Server {167 t.Helper()168 return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {169 conn, err := websocket.Accept(w, r, nil)170 if err != nil {171 t.Error(err)172 return173 }174 defer conn.CloseNow()175 for {176 _, raw, err := conn.Read(r.Context())177 if err != nil {178 return179 }180 var frame []json.RawMessage181 if json.Unmarshal(raw, &frame) != nil || len(frame) < 2 {182 continue183 }184 var kind, sub string185 _ = json.Unmarshal(frame[0], &kind)186 _ = json.Unmarshal(frame[1], &sub)187 switch kind {188 case "REQ":189 if sub == "tiny-agent" {190 _ = conn.Write(r.Context(), websocket.MessageText, []byte(`["EOSE","tiny-agent"]`))191 select {192 case ready <- struct{}{}:193 default:194 }195 continue196 }197 var filter map[string]any198 if len(frame) > 2 {199 _ = json.Unmarshal(frame[2], &filter)200 }201 var id string202 if value, ok := filter["ids"].([]any); ok && len(value) > 0 {203 id, _ = value[0].(string)204 }205 if e, ok := events[id]; ok {206 payload, _ := json.Marshal([]any{"EVENT", sub, e})207 _ = conn.Write(r.Context(), websocket.MessageText, payload)208 }209 if _, found := events[id]; found {210 payload, _ := json.Marshal([]any{"EOSE", sub})211 _ = conn.Write(r.Context(), websocket.MessageText, payload)212 }213 case "CLOSE":214 payload, _ := json.Marshal([]any{"CLOSED", sub, "subscription closed"})215 _ = conn.Write(r.Context(), websocket.MessageText, payload)216 if mu != nil && closed != nil {217 mu.Lock()218 *closed = append(*closed, sub)219 mu.Unlock()220 }221 }222 }223 }))224}