| 1 | package agentrunner |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "net/http" |
| 8 | "net/http/httptest" |
| 9 | "strings" |
| 10 | "sync" |
| 11 | "testing" |
| 12 | "time" |
| 13 | |
| 14 | "github.com/FelineStateMachine/tinyrelay/internal/event" |
| 15 | "github.com/coder/websocket" |
| 16 | ) |
| 17 | |
| 18 | const relayTestSecret = "0000000000000000000000000000000000000000000000000000000000000001" |
| 19 | |
| 20 | func 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.Mutex |
| 27 | var closed []string |
| 28 | ready := make(chan struct{}, 1) |
| 29 | server := relayTestServer(t, events, &mu, &closed, ready) |
| 30 | defer server.Close() |
| 31 | |
| 32 | client := RelayClient{URL: "ws" + strings.TrimPrefix(server.URL, "http" ), Secret: relayTestSecret, Rooms: []string{room}} |
| 33 | var dispatched int |
| 34 | 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 | } |
| 44 | |
| 45 | 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 | } |
| 77 | |
| 78 | func TestResolveRootRejectsInvalidQueriedEvents(t *testing.T) { |
| 79 | tests := []struct { |
| 80 | name string |
| 81 | make func(t *testing.T) (string, event.Event) |
| 82 | want string |
| 83 | }{ |
| 84 | {"wrong room" , func(t *testing.T) (string, event.Event) { |
| 85 | e := relayTestEvent(t, 11, "other" , "" ) |
| 86 | return "room" , e |
| 87 | }, "another room" }, |
| 88 | {"invalid kind" , func(t *testing.T) (string, event.Event) { |
| 89 | e := relayTestEvent(t, 1, "room" , "" ) |
| 90 | return "room" , e |
| 91 | }, "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 | } |
| 109 | |
| 110 | func 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 | } |
| 123 | |
| 124 | func 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 e |
| 135 | } |
| 136 | |
| 137 | func 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 != nil |
| 147 | client.mu.Unlock() |
| 148 | if connected { |
| 149 | break |
| 150 | } |
| 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 | } |
| 165 | |
| 166 | func 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 | return |
| 173 | } |
| 174 | defer conn.CloseNow() |
| 175 | for { |
| 176 | _, raw, err := conn.Read(r.Context()) |
| 177 | if err != nil { |
| 178 | return |
| 179 | } |
| 180 | var frame []json.RawMessage |
| 181 | if json.Unmarshal(raw, &frame) != nil || len(frame) < 2 { |
| 182 | continue |
| 183 | } |
| 184 | var kind, sub string |
| 185 | _ = 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 | continue |
| 196 | } |
| 197 | var filter map[string]any |
| 198 | if len(frame) > 2 { |
| 199 | _ = json.Unmarshal(frame[2], &filter) |
| 200 | } |
| 201 | var id string |
| 202 | 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 | } |