github.com/wallyworld/juju@v0.0.0-20161013125918-6cf1bc9d917a/worker/catacomb/fixture_test.go (about)

     1  // Copyright 2015 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package catacomb_test
     5  
     6  import (
     7  	"time"
     8  
     9  	jc "github.com/juju/testing/checkers"
    10  	gc "gopkg.in/check.v1"
    11  	"gopkg.in/tomb.v1"
    12  
    13  	coretesting "github.com/juju/juju/testing"
    14  	"github.com/juju/juju/worker"
    15  	"github.com/juju/juju/worker/catacomb"
    16  )
    17  
    18  type cleaner interface {
    19  	AddCleanup(func(*gc.C))
    20  }
    21  
    22  type fixture struct {
    23  	catacomb catacomb.Catacomb
    24  	cleaner  cleaner
    25  }
    26  
    27  func (fix *fixture) run(c *gc.C, task func(), init ...worker.Worker) error {
    28  	err := catacomb.Invoke(catacomb.Plan{
    29  		Site: &fix.catacomb,
    30  		Work: func() error { task(); return nil },
    31  		Init: init,
    32  	})
    33  	c.Assert(err, jc.ErrorIsNil)
    34  
    35  	select {
    36  	case <-fix.catacomb.Dead():
    37  	case <-time.After(coretesting.LongWait):
    38  		c.Fatalf("timed out")
    39  	}
    40  	return fix.catacomb.Wait()
    41  }
    42  
    43  func (fix *fixture) waitDying(c *gc.C) {
    44  	select {
    45  	case <-fix.catacomb.Dying():
    46  	case <-time.After(coretesting.LongWait):
    47  		c.Fatalf("timed out; still alive")
    48  	}
    49  }
    50  
    51  func (fix *fixture) assertDying(c *gc.C) {
    52  	select {
    53  	case <-fix.catacomb.Dying():
    54  	default:
    55  		c.Fatalf("still alive")
    56  	}
    57  }
    58  
    59  func (fix *fixture) assertNotDying(c *gc.C) {
    60  	select {
    61  	case <-fix.catacomb.Dying():
    62  		c.Fatalf("already dying")
    63  	default:
    64  	}
    65  }
    66  
    67  func (fix *fixture) assertDead(c *gc.C) {
    68  	select {
    69  	case <-fix.catacomb.Dead():
    70  	default:
    71  		c.Fatalf("not dead")
    72  	}
    73  }
    74  
    75  func (fix *fixture) assertNotDead(c *gc.C) {
    76  	select {
    77  	case <-fix.catacomb.Dead():
    78  		c.Fatalf("already dead")
    79  	default:
    80  	}
    81  }
    82  
    83  func (fix *fixture) assertAddAlive(c *gc.C, w *errorWorker) {
    84  	err := fix.catacomb.Add(w)
    85  	c.Assert(err, jc.ErrorIsNil)
    86  	w.waitStillAlive(c)
    87  }
    88  
    89  func (fix *fixture) startErrorWorker(c *gc.C, err error) *errorWorker {
    90  	ew := &errorWorker{}
    91  	go func() {
    92  		defer ew.tomb.Done()
    93  		defer ew.tomb.Kill(err)
    94  		<-ew.tomb.Dying()
    95  	}()
    96  	fix.cleaner.AddCleanup(func(_ *gc.C) {
    97  		ew.stop()
    98  	})
    99  	return ew
   100  }
   101  
   102  type errorWorker struct {
   103  	tomb tomb.Tomb
   104  }
   105  
   106  func (ew *errorWorker) Kill() {
   107  	ew.tomb.Kill(nil)
   108  }
   109  
   110  func (ew *errorWorker) Wait() error {
   111  	return ew.tomb.Wait()
   112  }
   113  
   114  func (ew *errorWorker) stop() {
   115  	ew.Kill()
   116  	ew.Wait()
   117  }
   118  
   119  func (ew *errorWorker) waitStillAlive(c *gc.C) {
   120  	select {
   121  	case <-ew.tomb.Dying():
   122  		c.Fatalf("already dying")
   123  	case <-time.After(coretesting.ShortWait):
   124  	}
   125  }
   126  
   127  func (ew *errorWorker) assertDead(c *gc.C) {
   128  	select {
   129  	case <-ew.tomb.Dead():
   130  	default:
   131  		c.Fatalf("not yet dead")
   132  	}
   133  }