github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/swarmkit/agent/reporter_test.go (about)

     1  package agent
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"math/rand"
     8  	"sync"
     9  	"testing"
    10  
    11  	"github.com/docker/swarmkit/api"
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  type uniqueStatus struct {
    16  	taskID string
    17  	status *api.TaskStatus
    18  }
    19  
    20  func TestReporter(t *testing.T) {
    21  	const ntasks = 100
    22  
    23  	var (
    24  		ctx      = context.Background()
    25  		statuses = make(map[string]*api.TaskStatus) // destination map
    26  		unique   = make(map[uniqueStatus]struct{})  // ensure we don't receive any status twice
    27  		mu       sync.Mutex
    28  		expected = make(map[string]*api.TaskStatus)
    29  		wg       sync.WaitGroup
    30  	)
    31  
    32  	reporter := newStatusReporter(ctx, statusReporterFunc(func(ctx context.Context, taskID string, status *api.TaskStatus) error {
    33  		if rand.Float64() > 0.9 {
    34  			return errors.New("status send failed")
    35  		}
    36  
    37  		mu.Lock()
    38  		defer mu.Unlock()
    39  
    40  		key := uniqueStatus{taskID, status}
    41  		// make sure we get the status only once.
    42  		if _, ok := unique[key]; ok {
    43  			t.Fatal("encountered status twice")
    44  		}
    45  
    46  		if status.State == api.TaskStateCompleted {
    47  			wg.Done()
    48  		}
    49  
    50  		unique[key] = struct{}{}
    51  		if current, ok := statuses[taskID]; ok {
    52  			if status.State <= current.State {
    53  				return nil // only allow forward updates
    54  			}
    55  		}
    56  
    57  		statuses[taskID] = status
    58  
    59  		return nil
    60  	}))
    61  
    62  	wg.Add(ntasks) // statuses will be reported!
    63  
    64  	for _, state := range []api.TaskState{
    65  		api.TaskStateAccepted,
    66  		api.TaskStatePreparing,
    67  		api.TaskStateReady,
    68  		api.TaskStateCompleted,
    69  	} {
    70  		for i := 0; i < ntasks; i++ {
    71  			taskID, status := fmt.Sprint(i), &api.TaskStatus{State: state}
    72  			expected[taskID] = status
    73  
    74  			// simulate pounding this with a bunch of goroutines
    75  			go func() {
    76  				if err := reporter.UpdateTaskStatus(ctx, taskID, status); err != nil {
    77  					assert.NoError(t, err, "sending should not fail")
    78  				}
    79  			}()
    80  
    81  		}
    82  	}
    83  
    84  	wg.Wait() // wait for the propagation
    85  	assert.NoError(t, reporter.Close())
    86  	mu.Lock()
    87  	defer mu.Unlock()
    88  
    89  	assert.Equal(t, expected, statuses)
    90  }