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  }