github.com/amp-space/amp-sdk-go@v0.7.6/stdlib/task/helpers_test.go (about)

     1  package task_test
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/amp-space/amp-sdk-go/stdlib/task"
    10  	"github.com/amp-space/amp-sdk-go/stdlib/testutils"
    11  )
    12  
    13  func makeItems() (item [5]*workItem) {
    14  	item[0] = &workItem{id: "a", processed: testutils.NewAwaiter()}
    15  	item[1] = &workItem{id: "b1", processed: testutils.NewAwaiter()}
    16  	item[2] = &workItem{id: "b2", processed: testutils.NewAwaiter()}
    17  	item[3] = &workItem{id: "c", processed: testutils.NewAwaiter()}
    18  	item[4] = &workItem{id: "d", processed: testutils.NewAwaiter()}
    19  	return
    20  }
    21  
    22  type workItem struct {
    23  	id        string
    24  	retry     bool
    25  	processed testutils.Awaiter
    26  	block     chan struct{}
    27  }
    28  
    29  func (i workItem) ID() task.PoolUniqueID { return i.id[:1] }
    30  
    31  func (i *workItem) Work(ctx context.Context) (retry bool) {
    32  	i.processed.ItHappened()
    33  	if i.block != nil {
    34  		select {
    35  		case <-i.block:
    36  			i.block = nil
    37  		case <-ctx.Done():
    38  			return false
    39  		}
    40  	}
    41  	retry = i.retry
    42  	i.retry = false
    43  	return retry
    44  }
    45  
    46  func requireHappened(t *testing.T, item *workItem, times int, wg *sync.WaitGroup) {
    47  	t.Helper()
    48  	defer wg.Done()
    49  	for i := 0; i < times; i++ {
    50  		item.processed.AwaitOrFail(t, 5*time.Second, item.id)
    51  	}
    52  	item.processed.NeverHappenedOrFail(t, 5*time.Second, item.id)
    53  }
    54  
    55  func requireNeverHappened(t *testing.T, item *workItem, wg *sync.WaitGroup) {
    56  	t.Helper()
    57  	defer wg.Done()
    58  	item.processed.NeverHappenedOrFail(t, 5*time.Second, item.id)
    59  }