43 lines
1.3 KiB
Go
43 lines
1.3 KiB
Go
package orbit
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
)
|
|
|
|
func TestCancellationNotificationAndExecutionBoundary(t *testing.T) {
|
|
r := workerTestRegistration(t)
|
|
r.CancellationNotifications = true
|
|
_, c := workerMock(t, r, func(ctx context.Context, c *websocket.Conn) {
|
|
body := workerTestBody(t, WorkerCancellation{TaskID: workerTestTask, AttemptID: workerTestAttempt, Epoch: 1, RequestedAt: time.Now().UTC()})
|
|
workerTestWrite(t, ctx, c, workerFrame{Type: "task.cancel", Body: body})
|
|
<-ctx.Done()
|
|
})
|
|
ctx, stop := context.WithTimeout(context.Background(), time.Second)
|
|
defer stop()
|
|
notice, err := c.NextCancellation(ctx)
|
|
if err != nil || notice.AttemptID != workerTestAttempt || notice.Epoch != 1 {
|
|
t.Fatalf("notice: %v", err)
|
|
}
|
|
cancelled, cancel := context.WithCancel(ctx)
|
|
cancel()
|
|
if _, err := c.NextCancellation(cancelled); !errors.Is(err, context.Canceled) {
|
|
t.Fatal("cancelled read did not end")
|
|
}
|
|
if c.Err() != nil {
|
|
t.Fatal("cancelled wait closed connection")
|
|
}
|
|
now := time.Now()
|
|
lease := WorkerLease{Active: true, State: "RUNNING", ServerTime: now, LeaseUntil: now.Add(time.Second), CancelRequestedAt: &now}
|
|
if err := lease.confirm(now); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, ok := lease.ExecutionDeadline(); ok {
|
|
t.Fatal("cancel intent still granted execution")
|
|
}
|
|
}
|