| 1 | package agentrunner |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "path/filepath" |
| 7 | "sync" |
| 8 | "testing" |
| 9 | "time" |
| 10 | ) |
| 11 | |
| 12 | func TestRunnerSerializesBatchesAndOtherRoomsProgress(t *testing.T) { |
| 13 | start := make(chan string, 3) |
| 14 | release := make(chan struct{}) |
| 15 | var mu sync.Mutex |
| 16 | var batches [][]Mention |
| 17 | r, err := New(Options{AllowedRooms: []string{"one" , "two" }, Handle: func(ctx context.Context, b []Mention) error { |
| 18 | mu.Lock() |
| 19 | batches = append(batches, b) |
| 20 | mu.Unlock() |
| 21 | start <- b[0].EventID |
| 22 | if b[0].EventID == "a" { |
| 23 | select { |
| 24 | case <-ctx.Done(): |
| 25 | return ctx.Err() |
| 26 | case <-release: |
| 27 | } |
| 28 | } |
| 29 | return nil |
| 30 | }}) |
| 31 | if err != nil { |
| 32 | t.Fatal(err) |
| 33 | } |
| 34 | r.Start(context.Background()) |
| 35 | defer r.Stop() |
| 36 | add := func(id, room string) { |
| 37 | t.Helper() |
| 38 | if err := r.Enqueue(Mention{EventID: id, Author: "alice" , Room: room}); err != nil { |
| 39 | t.Fatal(err) |
| 40 | } |
| 41 | } |
| 42 | add("a" , "one" ) |
| 43 | <-start |
| 44 | add("b" , "one" ) |
| 45 | add("c" , "one" ) |
| 46 | add("d" , "two" ) |
| 47 | select { |
| 48 | case id := <-start: |
| 49 | if id != "d" { |
| 50 | t.Fatal(id) |
| 51 | } |
| 52 | case <-time.After(time.Second): |
| 53 | t.Fatal("second room blocked" ) |
| 54 | } |
| 55 | if err := r.Enqueue(Mention{EventID: "b" , Author: "alice" , Room: "one" }); !errors.Is(err, ErrDuplicate) { |
| 56 | t.Fatal(err) |
| 57 | } |
| 58 | close(release) |
| 59 | ctx, cancel := context.WithTimeout(context.Background(), time.Second) |
| 60 | defer cancel() |
| 61 | if err := r.Drain(ctx); err != nil { |
| 62 | t.Fatal(err) |
| 63 | } |
| 64 | mu.Lock() |
| 65 | defer mu.Unlock() |
| 66 | if len(batches) != 3 || len(batches[2]) != 2 { |
| 67 | t.Fatalf("batches=%v" , batches) |
| 68 | } |
| 69 | } |
| 70 | func TestRunnerStopsActiveTaskAndRestoresJournal(t *testing.T) { |
| 71 | path := filepath.Join(t.TempDir(), "queue.json" ) |
| 72 | started := make(chan struct{}) |
| 73 | r, err := New(Options{AllowedRooms: []string{"one" }, StatePath: path, Handle: func(ctx context.Context, b []Mention) error { close(started); <-ctx.Done(); return ctx.Err() }}) |
| 74 | if err != nil { |
| 75 | t.Fatal(err) |
| 76 | } |
| 77 | r.Start(context.Background()) |
| 78 | r.Enqueue(Mention{Room: "one" , Author: "alice" , EventID: "a" }) |
| 79 | <-started |
| 80 | if r.CancelFrom("one" , "other" ) { |
| 81 | t.Fatal("another author canceled work" ) |
| 82 | } |
| 83 | done := make(chan struct{}) |
| 84 | go func() { r.Stop(); close(done) }() |
| 85 | select { |
| 86 | case <-done: |
| 87 | case <-time.After(time.Second): |
| 88 | t.Fatal("stop hung" ) |
| 89 | } |
| 90 | if err := r.Enqueue(Mention{Room: "one" , Author: "alice" , EventID: "b" }); !errors.Is(err, ErrStopped) { |
| 91 | t.Fatal(err) |
| 92 | } |
| 93 | got := make(chan string, 1) |
| 94 | restored, err := New(Options{AllowedRooms: []string{"one" }, StatePath: path, Handle: func(ctx context.Context, b []Mention) error { got <- b[0].EventID; return nil }}) |
| 95 | if err != nil { |
| 96 | t.Fatal(err) |
| 97 | } |
| 98 | restored.Start(context.Background()) |
| 99 | defer restored.Stop() |
| 100 | select { |
| 101 | case id := <-got: |
| 102 | if id != "a" { |
| 103 | t.Fatal(id) |
| 104 | } |
| 105 | case <-time.After(time.Second): |
| 106 | t.Fatal("unfinished work not recovered" ) |
| 107 | } |
| 108 | } |
| 109 | func TestRunnerRejectsUnconfiguredRoomsAndRecoversAfterFailure(t *testing.T) { |
| 110 | got := make(chan string, 2) |
| 111 | r, err := New(Options{AllowedRooms: []string{"one" }, QueueLimit: 1, Handle: func(ctx context.Context, b []Mention) error { |
| 112 | got <- b[0].EventID |
| 113 | return errors.New("process crashed" ) |
| 114 | }}) |
| 115 | if err != nil { |
| 116 | t.Fatal(err) |
| 117 | } |
| 118 | if err := r.Enqueue(Mention{Room: "two" , Author: "a" , EventID: "a" }); !errors.Is(err, ErrUnauthorizedRoom) { |
| 119 | t.Fatal(err) |
| 120 | } |
| 121 | r.Enqueue(Mention{Room: "one" , Author: "a" , EventID: "a" }) |
| 122 | if err := r.Enqueue(Mention{Room: "one" , Author: "a" , EventID: "b" }); !errors.Is(err, ErrQueueFull) { |
| 123 | t.Fatal(err) |
| 124 | } |
| 125 | r.Start(context.Background()) |
| 126 | defer r.Stop() |
| 127 | <-got |
| 128 | r.Enqueue(Mention{Room: "one" , Author: "a" , EventID: "b" }) |
| 129 | select { |
| 130 | case <-got: |
| 131 | case <-time.After(time.Second): |
| 132 | t.Fatal("queue died after process crash" ) |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | func TestRunnerDoesNotMixThreadBatches(t *testing.T) { |
| 137 | got := make(chan []Mention, 2) |
| 138 | r, err := New(Options{AllowedRooms: []string{"one" }, Handle: func(ctx context.Context, batch []Mention) error { |
| 139 | got <- batch |
| 140 | return nil |
| 141 | }}) |
| 142 | if err != nil { |
| 143 | t.Fatal(err) |
| 144 | } |
| 145 | r.Start(context.Background()) |
| 146 | defer r.Stop() |
| 147 | for _, m := range []Mention{{Room: "one" , Author: "a" , EventID: "a" , Root: "root-a" }, {Room: "one" , Author: "b" , EventID: "b" , Root: "root-a" }, {Room: "one" , Author: "c" , EventID: "c" , Root: "root-b" }} { |
| 148 | if err := r.Enqueue(m); err != nil { |
| 149 | t.Fatal(err) |
| 150 | } |
| 151 | } |
| 152 | first := <-got |
| 153 | second := <-got |
| 154 | if len(first) != 2 || first[0].Root != "root-a" || len(second) != 1 || second[0].Root != "root-b" { |
| 155 | t.Fatalf("batches=%v,%v" , first, second) |
| 156 | } |
| 157 | } |
| 158 | |
| 159 | func TestRunnerCancelFromKeepsOtherAuthorsQueued(t *testing.T) { |
| 160 | started := make(chan struct{}) |
| 161 | release := make(chan struct{}) |
| 162 | var startOnce sync.Once |
| 163 | r, err := New(Options{AllowedRooms: []string{"one" }, Handle: func(ctx context.Context, batch []Mention) error { |
| 164 | startOnce.Do(func() { close(started) }) |
| 165 | select { |
| 166 | case <-release: |
| 167 | return nil |
| 168 | case <-ctx.Done(): |
| 169 | return ctx.Err() |
| 170 | } |
| 171 | }}) |
| 172 | if err != nil { |
| 173 | t.Fatal(err) |
| 174 | } |
| 175 | r.Start(context.Background()) |
| 176 | defer r.Stop() |
| 177 | if err := r.Enqueue(Mention{Room: "one" , Author: "alice" , EventID: "a" , Root: "root" }); err != nil { |
| 178 | t.Fatal(err) |
| 179 | } |
| 180 | <-started |
| 181 | if err := r.Enqueue(Mention{Room: "one" , Author: "bob" , EventID: "b" , Root: "root-b" }); err != nil { |
| 182 | t.Fatal(err) |
| 183 | } |
| 184 | if !r.CancelFrom("one" , "alice" ) { |
| 185 | t.Fatal("alice cancellation was not accepted" ) |
| 186 | } |
| 187 | close(release) |
| 188 | ctx, cancel := context.WithTimeout(context.Background(), time.Second) |
| 189 | defer cancel() |
| 190 | if err := r.Drain(ctx); err != nil { |
| 191 | t.Fatal(err) |
| 192 | } |
| 193 | if err := r.Enqueue(Mention{Room: "one" , Author: "bob" , EventID: "b" }); !errors.Is(err, ErrDuplicate) { |
| 194 | t.Fatalf("bob event should remain seen: %v" , err) |
| 195 | } |
| 196 | } |