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