launchpad.net/~rogpeppe/juju-core/500-errgo-fix@v0.0.0-20140213181702-000000002356/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  	"launchpad.net/juju-core/provider/common"
     8  	"launchpad.net/juju-core/utils"
     9  )
    10  
    11  // impatientAttempt is an extremely short polling time suitable for tests.
    12  // It polls at least once, never delays, and times out very quickly.
    13  var impatientAttempt = utils.AttemptStrategy{}
    14  
    15  // savedAttemptStrategy holds the state needed to restore an AttemptStrategy's
    16  // original setting.
    17  type savedAttemptStrategy struct {
    18  	address  *utils.AttemptStrategy
    19  	original utils.AttemptStrategy
    20  }
    21  
    22  // saveAttemptStrategies captures the information required to restore the
    23  // given AttemptStrategy objects.
    24  func saveAttemptStrategies(strategies []*utils.AttemptStrategy) []savedAttemptStrategy {
    25  	savedStrategies := make([]savedAttemptStrategy, len(strategies))
    26  	for index, strategy := range strategies {
    27  		savedStrategies[index] = savedAttemptStrategy{
    28  			address:  strategy,
    29  			original: *strategy,
    30  		}
    31  	}
    32  	return savedStrategies
    33  }
    34  
    35  // restore returns a saved strategy to its original configuration.
    36  func (saved savedAttemptStrategy) restore() {
    37  	*saved.address = saved.original
    38  }
    39  
    40  // restoreAttemptStrategies restores multiple saved AttemptStrategies.
    41  func restoreAttemptStrategies(strategies []savedAttemptStrategy) {
    42  	for _, saved := range strategies {
    43  		saved.restore()
    44  	}
    45  }
    46  
    47  // internalPatchAttemptStrategies sets the given AttemptStrategy objects
    48  // to the impatientAttempt configuration, and returns a function that restores
    49  // the original configurations.
    50  func internalPatchAttemptStrategies(strategies []*utils.AttemptStrategy) func() {
    51  	snapshot := saveAttemptStrategies(strategies)
    52  	for _, strategy := range strategies {
    53  		*strategy = impatientAttempt
    54  	}
    55  	return func() { restoreAttemptStrategies(snapshot) }
    56  }
    57  
    58  // TODO: Everything up to this point is generic.  Move it to utils?
    59  
    60  // PatchAttemptStrategies patches environs' global polling strategy, plus any
    61  // otther AttemptStrategy objects whose addresses you pass, to very short
    62  // polling and timeout times so that tests can run fast.
    63  // It returns a cleanup function that restores the original settings.  You must
    64  // call this afterwards.
    65  func PatchAttemptStrategies(strategies ...*utils.AttemptStrategy) func() {
    66  	// The one irregularity here is that LongAttempt goes on the list of
    67  	// strategies that need patching.  To keep testing simple, we treat
    68  	// the given attempts and LongAttempt as a single slice from here on.
    69  	combinedStrategies := append(strategies, &common.LongAttempt, &common.ShortAttempt)
    70  	return internalPatchAttemptStrategies(combinedStrategies)
    71  }