github.com/cloud-green/juju@v0.0.0-20151002100041-a00291338d3d/worker/periodicworker_test.go (about)

     1  // Copyright 2014 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package worker
     5  
     6  import (
     7  	"time"
     8  
     9  	gc "gopkg.in/check.v1"
    10  
    11  	"github.com/juju/juju/testing"
    12  )
    13  
    14  type periodicWorkerSuite struct {
    15  	testing.BaseSuite
    16  }
    17  
    18  var (
    19  	_                   = gc.Suite(&periodicWorkerSuite{})
    20  	defaultPeriod       = time.Second
    21  	defaultFireOnceWait = defaultPeriod / 2
    22  )
    23  
    24  func (s *periodicWorkerSuite) TestWait(c *gc.C) {
    25  	funcHasRun := make(chan struct{})
    26  	doWork := func(_ <-chan struct{}) error {
    27  		funcHasRun <- struct{}{}
    28  		return testError
    29  	}
    30  
    31  	w := NewPeriodicWorker(doWork, defaultPeriod, NewTimer)
    32  	defer func() { c.Assert(Stop(w), gc.Equals, testError) }()
    33  	select {
    34  	case <-funcHasRun:
    35  	case <-time.After(testing.ShortWait):
    36  		c.Fatalf("The doWork function should have been called by now")
    37  	}
    38  	w.Kill()
    39  	c.Assert(w.Wait(), gc.Equals, testError)
    40  	select {
    41  	case <-funcHasRun:
    42  		c.Fatalf("After the kill we don't expect anymore calls to the function")
    43  	case <-time.After(defaultFireOnceWait):
    44  	}
    45  }
    46  
    47  // TestWaitNil starts a periodicWorker asserts that after
    48  // killing the worker Wait() returns nil after at least
    49  // one call of the doWork function
    50  func (s *periodicWorkerSuite) TestWaitNil(c *gc.C) {
    51  	funcHasRun := make(chan struct{})
    52  	doWork := func(_ <-chan struct{}) error {
    53  		funcHasRun <- struct{}{}
    54  		return nil
    55  	}
    56  
    57  	w := NewPeriodicWorker(doWork, defaultPeriod, NewTimer)
    58  	defer func() { c.Assert(Stop(w), gc.IsNil) }()
    59  	select {
    60  	case <-funcHasRun:
    61  	case <-time.After(defaultFireOnceWait):
    62  		c.Fatalf("The doWork function should have been called by now")
    63  	}
    64  	w.Kill()
    65  	c.Assert(w.Wait(), gc.Equals, nil)
    66  }
    67  
    68  // TestKill starts a periodic worker and Kills it
    69  // it expects the doWork function to be notified of this and the error from
    70  // doWork is returned by Wait()
    71  func (s *periodicWorkerSuite) TestKill(c *gc.C) {
    72  	tests := []struct {
    73  		ReturnValue   error
    74  		ExpectedValue error
    75  	}{
    76  		{nil, nil},
    77  		{testError, testError},
    78  		{ErrKilled, nil},
    79  	}
    80  
    81  	for i, test := range tests {
    82  		c.Logf("Running test %d\n", i)
    83  		runKillTest(c, test.ReturnValue, test.ExpectedValue)
    84  	}
    85  }
    86  
    87  func runKillTest(c *gc.C, returnValue, expected error) {
    88  	ready := make(chan struct{})
    89  	doWorkNotification := make(chan struct{})
    90  	doWork := func(stopCh <-chan struct{}) error {
    91  		close(ready)
    92  		<-stopCh
    93  		close(doWorkNotification)
    94  		return returnValue
    95  	}
    96  
    97  	w := NewPeriodicWorker(doWork, defaultPeriod, NewTimer)
    98  	defer func() { c.Assert(Stop(w), gc.Equals, expected) }()
    99  
   100  	select {
   101  	case <-ready:
   102  	case <-time.After(testing.LongWait):
   103  		c.Fatalf("The doWork call should be ready by now")
   104  	}
   105  	w.Kill()
   106  	select {
   107  	case <-doWorkNotification:
   108  	case <-time.After(testing.LongWait):
   109  		c.Fatalf("The doWork function should have been notified of the stop by now")
   110  	}
   111  	c.Assert(w.Wait(), gc.Equals, expected)
   112  
   113  	// test we can kill again without a panic and our death reason stays intact
   114  	w.Kill()
   115  }
   116  
   117  // TestCallUntilKilled checks that our function is called
   118  // at least 5 times, and that with a period of 500ms each call is made
   119  // in a reasonable time
   120  func (s *periodicWorkerSuite) TestCallUntilKilled(c *gc.C) {
   121  	funcHasRun := make(chan struct{}, 5)
   122  	doWork := func(_ <-chan struct{}) error {
   123  		funcHasRun <- struct{}{}
   124  		return nil
   125  	}
   126  
   127  	period := time.Millisecond * 500
   128  	unacceptableWait := time.Second * 10
   129  	w := NewPeriodicWorker(doWork, period, NewTimer)
   130  	defer func() { c.Assert(Stop(w), gc.IsNil) }()
   131  	for i := 0; i < 5; i++ {
   132  		select {
   133  		case <-funcHasRun:
   134  		case <-time.After(unacceptableWait):
   135  			c.Fatalf("The function should have been called again by now")
   136  		}
   137  	}
   138  	w.Kill()
   139  	c.Assert(w.Wait(), gc.Equals, nil)
   140  }