github.com/grailbio/base@v0.0.11/sync/workerpool/workerpool_test.go (about)

     1  package workerpool_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync/atomic"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/grailbio/base/grail"
    11  	"github.com/grailbio/base/sync/workerpool"
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  // TODO(pknudsgaard): Refactor to table-driven tests?
    16  
    17  func TestNoTasks(t *testing.T) {
    18  	ctx := context.Background()
    19  
    20  	wp := workerpool.New(ctx, 10)
    21  	wp.Wait()
    22  }
    23  
    24  type WaitTask struct {
    25  	delay     time.Duration
    26  	Completed int64
    27  }
    28  
    29  func (wt *WaitTask) Do(ctx *workerpool.TaskGroup) error {
    30  	time.Sleep(wt.delay)
    31  	atomic.AddInt64(&wt.Completed, 1)
    32  
    33  	return nil
    34  }
    35  
    36  func TestSingleTaskBlock(t *testing.T) {
    37  	wp := workerpool.New(context.Background(), 10)
    38  	ctx := wp.NewTaskGroup("test", nil)
    39  	wt := WaitTask{
    40  		delay: 100 * time.Millisecond,
    41  	}
    42  	assert.Equal(t, true, ctx.Enqueue(&wt, true))
    43  	ctx.Wait()
    44  	assert.Equal(t, int64(1), wt.Completed)
    45  	wp.Wait()
    46  }
    47  
    48  func TestSingleTaskNoBlock(t *testing.T) {
    49  	wp := workerpool.New(context.Background(), 10)
    50  	ctx := wp.NewTaskGroup("test", nil)
    51  	wt := WaitTask{
    52  		delay: 100 * time.Millisecond,
    53  	}
    54  	assert.Equal(t, true, ctx.Enqueue(&wt, false))
    55  	ctx.Wait()
    56  	assert.Equal(t, int64(1), wt.Completed)
    57  	wp.Wait()
    58  }
    59  
    60  func TestContextManyTasks(t *testing.T) {
    61  	wp := workerpool.New(context.Background(), 10)
    62  	ctx := wp.NewTaskGroup("test", nil)
    63  	wt := WaitTask{
    64  		delay: 10 * time.Millisecond,
    65  	}
    66  	for i := 0; i < 1000; i++ {
    67  		assert.Equal(t, true, ctx.Enqueue(&wt, true))
    68  	}
    69  	ctx.Wait()
    70  	assert.Equal(t, int64(1000), wt.Completed)
    71  	wp.Wait()
    72  }
    73  
    74  func TestContextTimeout(t *testing.T) {
    75  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
    76  	defer cancel()
    77  
    78  	wp := workerpool.New(ctx, 10)
    79  	wpctx := wp.NewTaskGroup("test", nil)
    80  	wt := WaitTask{
    81  		delay: time.Second,
    82  	}
    83  	for i := 0; i < 100; i++ {
    84  		assert.Equal(t, true, wpctx.Enqueue(&wt, true))
    85  	}
    86  	wpctx.Wait()
    87  	assert.True(t, wt.Completed > 0)
    88  	assert.True(t, wt.Completed < 100, fmt.Sprintf("All the tasks Completed %v", wt.Completed))
    89  	assert.Equal(t, "context deadline exceeded", ctx.Err().Error())
    90  	wp.Wait()
    91  }
    92  
    93  func TestMain(m *testing.M) {
    94  	shutdown := grail.Init()
    95  	defer shutdown()
    96  	m.Run()
    97  }