github.com/mailgun/holster/v4@v4.20.0/syncutil/waitgroup_test.go (about)

     1  /*
     2  Copyright 2017 Mailgun Technologies Inc
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8  	http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  package syncutil_test
    17  
    18  import (
    19  	"sync/atomic"
    20  	"testing"
    21  	"time"
    22  
    23  	linq "github.com/ahmetb/go-linq"
    24  	"github.com/mailgun/holster/v4/errors"
    25  	"github.com/mailgun/holster/v4/syncutil"
    26  	"github.com/stretchr/testify/suite"
    27  )
    28  
    29  type WaitGroupTestSuite struct {
    30  	suite.Suite
    31  }
    32  
    33  func TestWaitGroup(t *testing.T) {
    34  	suite.Run(t, new(WaitGroupTestSuite))
    35  }
    36  
    37  func (s *WaitGroupTestSuite) TestRun() {
    38  	var wg syncutil.WaitGroup
    39  
    40  	items := []error{
    41  		errors.New("Error 1"),
    42  		errors.New("Error 2"),
    43  	}
    44  
    45  	// Iterate over a thing and doing some long running thing for each
    46  	for _, item := range items {
    47  		wg.Run(func(item interface{}) error {
    48  			// Do some long running thing
    49  			time.Sleep(time.Nanosecond * 50)
    50  			// Return an error for testing
    51  			return item.(error)
    52  		}, item)
    53  	}
    54  
    55  	errs := wg.Wait()
    56  	s.NotNil(errs)
    57  	s.Equal(2, len(errs))
    58  	s.Equal(true, linq.From(errs).Contains(items[0]))
    59  	s.Equal(true, linq.From(errs).Contains(items[1]))
    60  }
    61  
    62  func (s *WaitGroupTestSuite) TestGo() {
    63  	var wg syncutil.WaitGroup
    64  	result := make(chan struct{})
    65  
    66  	wg.Go(func() {
    67  		// Do some long running thing
    68  		time.Sleep(time.Nanosecond * 500)
    69  		result <- struct{}{}
    70  	})
    71  
    72  	wg.Go(func() {
    73  		// Do some long running thing
    74  		time.Sleep(time.Nanosecond * 50)
    75  		result <- struct{}{}
    76  	})
    77  
    78  OUT:
    79  	for i := 0; i < 2; {
    80  		select {
    81  		case <-result:
    82  			i++
    83  		case <-time.After(time.Second):
    84  			s.Fail("waited to long for Go() to run")
    85  			break OUT
    86  		}
    87  	}
    88  
    89  	errs := wg.Wait()
    90  	s.Nil(errs)
    91  }
    92  
    93  func (s *WaitGroupTestSuite) TestLoop() {
    94  	pipe := make(chan int32)
    95  	var wg syncutil.WaitGroup
    96  	var count int32
    97  
    98  	wg.Loop(func() bool {
    99  		inc, ok := <-pipe
   100  		if !ok {
   101  			return false
   102  		}
   103  		atomic.AddInt32(&count, inc)
   104  		return true
   105  	})
   106  
   107  	// Feed the loop some numbers and close the pipe
   108  	pipe <- 1
   109  	pipe <- 5
   110  	pipe <- 10
   111  	close(pipe)
   112  
   113  	// Wait for the routine to end
   114  	// no error collection when using Loop()
   115  	errs := wg.Wait()
   116  	s.Nil(errs)
   117  	s.Equal(int32(16), count)
   118  }
   119  
   120  func (s *WaitGroupTestSuite) TestUntil() {
   121  	pipe := make(chan int32)
   122  	var wg syncutil.WaitGroup
   123  	var count int32
   124  
   125  	wg.Until(func(done chan struct{}) bool {
   126  		select {
   127  		case inc := <-pipe:
   128  			atomic.AddInt32(&count, inc)
   129  		case <-done:
   130  			return false
   131  		}
   132  		return true
   133  	})
   134  
   135  	wg.Until(func(done chan struct{}) bool {
   136  		select {
   137  		case inc := <-pipe:
   138  			atomic.AddInt32(&count, inc)
   139  		case <-done:
   140  			return false
   141  		}
   142  		return true
   143  	})
   144  
   145  	// Feed the loop some numbers and close the pipe
   146  	pipe <- 1
   147  	pipe <- 5
   148  	pipe <- 10
   149  
   150  	// Wait for the routine to end
   151  	wg.Stop()
   152  	s.Equal(int32(16), count)
   153  }