github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/swarmkit/agent/task_test.go (about) 1 package agent 2 3 import ( 4 "context" 5 "testing" 6 "time" 7 8 "github.com/docker/swarmkit/agent/exec" 9 "github.com/docker/swarmkit/api" 10 "github.com/sirupsen/logrus" 11 "github.com/stretchr/testify/assert" 12 ) 13 14 func init() { 15 logrus.SetLevel(logrus.DebugLevel) 16 } 17 18 func TestTaskManager(t *testing.T) { 19 ctx := context.Background() 20 task := &api.Task{ 21 Status: api.TaskStatus{}, 22 DesiredState: api.TaskStateAccepted, 23 } 24 accepted := make(chan struct{}) 25 ready := make(chan struct{}) 26 shutdown := make(chan struct{}) 27 ctlr := &controllerStub{t: t, calls: map[string]int{}} 28 29 tm := newTaskManager(ctx, task, ctlr, statusReporterFunc(func(ctx context.Context, taskID string, status *api.TaskStatus) error { 30 switch status.State { 31 case api.TaskStateAccepted: 32 select { 33 case <-accepted: 34 default: 35 close(accepted) 36 } 37 case api.TaskStatePreparing: 38 case api.TaskStateReady: 39 select { 40 case <-ready: 41 default: 42 close(ready) 43 } 44 case api.TaskStateStarting: 45 case api.TaskStateRunning: 46 select { 47 case <-ready: 48 default: 49 t.Fatal("should be running before ready") 50 } 51 case api.TaskStateCompleted: 52 select { 53 case <-shutdown: 54 default: 55 close(shutdown) 56 } 57 default: 58 t.Fatalf("unexpected state encountered: %v", status.State) 59 } 60 61 return nil 62 })) 63 64 acceptedWait := accepted 65 readyWait := ready 66 shutdownWait := shutdown 67 for { 68 select { 69 case <-acceptedWait: 70 task.DesiredState = api.TaskStateReady // proceed to ready 71 assert.NoError(t, tm.Update(ctx, task)) 72 acceptedWait = nil 73 case <-readyWait: 74 time.Sleep(time.Second) 75 task.DesiredState = api.TaskStateRunning // proceed to running. 76 assert.NoError(t, tm.Update(ctx, task)) 77 readyWait = nil 78 case <-shutdownWait: 79 assert.NoError(t, tm.Close()) 80 select { 81 case <-tm.closed: 82 default: 83 t.Fatal("not actually closed") 84 } 85 86 assert.NoError(t, tm.Close()) // hit a second time to make sure it behaves 87 assert.Equal(t, tm.Update(ctx, task), ErrClosed) 88 89 assert.Equal(t, map[string]int{ 90 "start": 1, 91 "wait": 1, 92 "prepare": 1, 93 "update": 2}, ctlr.calls) 94 return 95 case <-ctx.Done(): 96 t.Fatal(ctx.Err()) 97 } 98 } 99 } 100 101 type controllerStub struct { 102 t *testing.T 103 exec.Controller 104 105 calls map[string]int 106 } 107 108 func (cs *controllerStub) Prepare(ctx context.Context) error { 109 cs.calls["prepare"]++ 110 cs.t.Log("(*controllerStub).Prepare") 111 return nil 112 } 113 114 func (cs *controllerStub) Start(ctx context.Context) error { 115 cs.calls["start"]++ 116 cs.t.Log("(*controllerStub).Start") 117 return nil 118 } 119 120 func (cs *controllerStub) Wait(ctx context.Context) error { 121 cs.calls["wait"]++ 122 cs.t.Log("(*controllerStub).Wait") 123 return nil 124 } 125 126 func (cs *controllerStub) Update(ctx context.Context, task *api.Task) error { 127 cs.calls["update"]++ 128 cs.t.Log("(*controllerStub).Update") 129 return nil 130 } 131 132 func (cs *controllerStub) Remove(ctx context.Context) error { 133 cs.calls["remove"]++ 134 cs.t.Log("(*controllerStub).Remove") 135 return nil 136 } 137 138 func (cs *controllerStub) Close() error { 139 cs.calls["close"]++ 140 cs.t.Log("(*controllerStub).Close") 141 return nil 142 }