github.com/ngicks/gokugen@v0.0.5/scheduler/scheduler_test.go (about)

     1  package scheduler_test
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"sync/atomic"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/ngicks/gokugen/scheduler"
    11  	"go.uber.org/goleak"
    12  )
    13  
    14  func TestMain(m *testing.M) {
    15  	goleak.VerifyTestMain(m)
    16  }
    17  func TestScheduler(t *testing.T) {
    18  	t.Run("scheduler basic usage", func(t *testing.T) {
    19  		// TODO: change test or internal structure to remove flakiness.
    20  		workerNum := int64(5)
    21  		s := scheduler.NewScheduler(uint(workerNum), 0)
    22  		now := time.Now()
    23  
    24  		taskTicker := make(chan struct{}, 5)
    25  		var count uint32
    26  		taskFunc := func(taskCtx context.Context, t time.Time) {
    27  			atomic.AddUint32(&count, 1)
    28  			<-taskTicker
    29  		}
    30  
    31  		for i := 0; i < 10; i++ {
    32  			s.Schedule(scheduler.NewTask(now.Add(time.Millisecond), taskFunc))
    33  		}
    34  
    35  		ctx, cancel := context.WithCancel(context.Background())
    36  		ctxChangeCh := make(chan struct{})
    37  		go func() {
    38  			<-ctxChangeCh
    39  			s.Start(ctx)
    40  		}()
    41  		ctxChangeCh <- struct{}{}
    42  
    43  		if c := atomic.LoadUint32(&count); c != 0 {
    44  			t.Fatalf("worker number limitation is not wokring: should be %d, but %d", 0, c)
    45  		}
    46  		if s.ActiveWorkerNum() != 0 {
    47  			t.Fatalf("there must no be any active worker.")
    48  		}
    49  		time.Sleep(2 * time.Millisecond)
    50  		if c := atomic.LoadUint32(&count); c != 5 {
    51  			t.Fatalf("worker number limitation is not wokring: should be %d, but %d", 5, c)
    52  		}
    53  		if s.ActiveWorkerNum() != workerNum {
    54  			t.Fatalf("Active worker must be %d, but %d", workerNum, s.ActiveWorkerNum())
    55  		}
    56  		for i := 0; i < 3; i++ {
    57  			taskTicker <- struct{}{}
    58  		}
    59  		time.Sleep(time.Millisecond)
    60  		if c := atomic.LoadUint32(&count); c != 8 {
    61  			t.Fatalf("worker number limitation is not wokring: should be %d, but %d", 8, c)
    62  		}
    63  
    64  		for i := 0; i < 5; i++ {
    65  			taskTicker <- struct{}{}
    66  		}
    67  		time.Sleep(time.Millisecond)
    68  		if activeWorker := s.ActiveWorkerNum(); activeWorker != 2 {
    69  			t.Fatalf("Active worker must be %d, but %d", 2, activeWorker)
    70  		}
    71  
    72  		for i := 0; i < 2; i++ {
    73  			taskTicker <- struct{}{}
    74  		}
    75  		time.Sleep(time.Millisecond)
    76  		if c := atomic.LoadUint32(&count); c != 10 {
    77  			t.Fatalf("worker number limitation is not wokring: should be %d, but %d", 10, c)
    78  		}
    79  		cancel()
    80  
    81  		s.End()
    82  	})
    83  
    84  	t.Run("ended scheduler can not be started again.", func(t *testing.T) {
    85  		s := scheduler.NewScheduler(5, 0)
    86  		s.End()
    87  
    88  		if err := s.Start(context.TODO()); err != scheduler.ErrAlreadyEnded {
    89  			t.Fatalf("mismatched error; must be %v, but is %v", scheduler.ErrAlreadyEnded, err)
    90  		}
    91  	})
    92  
    93  	t.Run("restart", func(t *testing.T) {
    94  		s := scheduler.NewScheduler(5, 0)
    95  
    96  		now := time.Now()
    97  		var count uint32
    98  		taskFunc := func(taskCtx context.Context, t time.Time) {
    99  			atomic.AddUint32(&count, 1)
   100  		}
   101  
   102  		s.Schedule(scheduler.NewTask(now.Add(time.Millisecond), taskFunc))
   103  		s.Schedule(scheduler.NewTask(now.Add(10*time.Millisecond), taskFunc))
   104  
   105  		wg := sync.WaitGroup{}
   106  		ctx, cancel := context.WithCancel(context.Background())
   107  
   108  		ctxChangeCh := make(chan struct{})
   109  		wg.Add(1)
   110  		go func() {
   111  			<-ctxChangeCh
   112  			s.Start(ctx)
   113  			wg.Done()
   114  		}()
   115  		ctxChangeCh <- struct{}{}
   116  
   117  		time.Sleep(5 * time.Millisecond)
   118  		if c := atomic.LoadUint32(&count); c != 1 {
   119  			t.Fatalf("worker number limitation is not wokring: should be %d, but %d", 1, c)
   120  		}
   121  		cancel()
   122  
   123  		wg.Wait()
   124  
   125  		ctx, cancel = context.WithCancel(context.Background())
   126  		wg.Add(1)
   127  		var err error
   128  		go func() {
   129  			<-ctxChangeCh
   130  			err = s.Start(ctx)
   131  			wg.Done()
   132  		}()
   133  		ctxChangeCh <- struct{}{}
   134  
   135  		time.Sleep(15 * time.Millisecond)
   136  		if c := atomic.LoadUint32(&count); c != 2 {
   137  			t.Fatalf("worker number limitation is not wokring: should be %d, but %d", 2, c)
   138  		}
   139  
   140  		cancel()
   141  		wg.Wait()
   142  
   143  		if err != nil {
   144  			t.Fatalf("could not restart; %v", err)
   145  		}
   146  
   147  		s.End()
   148  	})
   149  
   150  }