weichenandCursor
ea2b9f121e
Allow tests and shutdown to drain in-memory workers before TempDir cleanup. Co-authored-by: Cursor <cursoragent@cursor.com>
63 lines
1.4 KiB
Go
63 lines
1.4 KiB
Go
package taskrun
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/bizdev"
|
|
"chat-agent-system/internal/stream"
|
|
)
|
|
|
|
func TestManager_WaitIdle_WaitsForWorker(t *testing.T) {
|
|
m := NewManager(nil, nil)
|
|
block := make(chan struct{})
|
|
m.SetExecutor(
|
|
func(ctx context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
|
|
select {
|
|
case <-ctx.Done():
|
|
return &bizdev.TaskRunResult{Success: false}, ctx.Err()
|
|
case <-block:
|
|
return &bizdev.TaskRunResult{Success: true}, nil
|
|
}
|
|
},
|
|
nil,
|
|
nil,
|
|
)
|
|
|
|
rec := &RunRecord{
|
|
RunID: "run-wait-idle",
|
|
SessionID: "sess",
|
|
UserID: 1,
|
|
Status: "pending",
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
if _, err := m.CreateAndStart(context.Background(), rec.RunID, rec, bizdev.TaskRunRequest{}, bizdev.StageRunRequest{}, ""); err != nil {
|
|
t.Fatalf("CreateAndStart: %v", err)
|
|
}
|
|
|
|
idleDone := make(chan error, 1)
|
|
go func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
idleDone <- m.WaitIdle(ctx)
|
|
}()
|
|
|
|
select {
|
|
case err := <-idleDone:
|
|
t.Fatalf("WaitIdle returned before worker released: %v", err)
|
|
case <-time.After(50 * time.Millisecond):
|
|
}
|
|
|
|
close(block)
|
|
select {
|
|
case err := <-idleDone:
|
|
if err != nil {
|
|
t.Fatalf("WaitIdle: %v", err)
|
|
}
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("WaitIdle timed out after releasing worker")
|
|
}
|
|
}
|