github.com/wallyworld/juju@v0.0.0-20161013125918-6cf1bc9d917a/environs/testing/polling.go (about)

     1  // Copyright 2012, 2013 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package testing
     5  
     6  import (
     7  	"github.com/juju/utils"
     8  
     9  	"github.com/juju/juju/environs"
    10  	"github.com/juju/juju/provider/common"
    11  )
    12  
    13  // impatientAttempt is an extremely short polling time suitable for tests.
    14  // It polls at least once, never delays, and times out very quickly.
    15  //
    16  // TODO(katco): 2016-08-09: lp:1611427
    17  var impatientAttempt = utils.AttemptStrategy{}
    18  
    19  // savedAttemptStrategy holds the state needed to restore an AttemptStrategy's
    20  // original setting.
    21  //
    22  // TODO(katco): 2016-08-09: lp:1611427
    23  type savedAttemptStrategy struct {
    24  	address  *utils.AttemptStrategy
    25  	original utils.AttemptStrategy
    26  }
    27  
    28  // saveAttemptStrategies captures the information required to restore the
    29  // given AttemptStrategy objects.
    30  //
    31  // TODO(katco): 2016-08-09: lp:1611427
    32  func saveAttemptStrategies(strategies []*utils.AttemptStrategy) []savedAttemptStrategy {
    33  	savedStrategies := make([]savedAttemptStrategy, len(strategies))
    34  	for index, strategy := range strategies {
    35  		savedStrategies[index] = savedAttemptStrategy{
    36  			address:  strategy,
    37  			original: *strategy,
    38  		}
    39  	}
    40  	return savedStrategies
    41  }
    42  
    43  // restore returns a saved strategy to its original configuration.
    44  func (saved savedAttemptStrategy) restore() {
    45  	*saved.address = saved.original
    46  }
    47  
    48  // restoreAttemptStrategies restores multiple saved AttemptStrategies.
    49  func restoreAttemptStrategies(strategies []savedAttemptStrategy) {
    50  	for _, saved := range strategies {
    51  		saved.restore()
    52  	}
    53  }
    54  
    55  // internalPatchAttemptStrategies sets the given AttemptStrategy objects
    56  // to the impatientAttempt configuration, and returns a function that restores
    57  // the original configurations.
    58  //
    59  // TODO(katco): 2016-08-09: lp:1611427
    60  func internalPatchAttemptStrategies(strategies []*utils.AttemptStrategy) func() {
    61  	snapshot := saveAttemptStrategies(strategies)
    62  	for _, strategy := range strategies {
    63  		*strategy = impatientAttempt
    64  	}
    65  	return func() { restoreAttemptStrategies(snapshot) }
    66  }
    67  
    68  // TODO: Everything up to this point is generic.  Move it to utils?
    69  
    70  // PatchAttemptStrategies patches environs' global polling strategy, plus any
    71  // otther AttemptStrategy objects whose addresses you pass, to very short
    72  // polling and timeout times so that tests can run fast.
    73  // It returns a cleanup function that restores the original settings.  You must
    74  // call this afterwards.
    75  //
    76  // TODO(katco): 2016-08-09: lp:1611427
    77  func PatchAttemptStrategies(strategies ...*utils.AttemptStrategy) func() {
    78  	// The one irregularity here is that LongAttempt goes on the list of
    79  	// strategies that need patching.  To keep testing simple, we treat
    80  	// the given attempts and LongAttempt as a single slice from here on.
    81  	combinedStrategies := append(
    82  		strategies,
    83  		&common.LongAttempt,
    84  		&common.ShortAttempt,
    85  		&environs.AddressesRefreshAttempt,
    86  	)
    87  	return internalPatchAttemptStrategies(combinedStrategies)
    88  }