github.com/cloudbase/juju-core@v0.0.0-20140504232958-a7271ac7912f/testing/testbase/cleanup.go (about) 1 // Copyright 2013 Canonical Ltd. 2 // Licensed under the AGPLv3, see LICENCE file for details. 3 4 package testbase 5 6 import ( 7 "os/exec" 8 9 gc "launchpad.net/gocheck" 10 ) 11 12 type CleanupFunc func(*gc.C) 13 type cleanupStack []CleanupFunc 14 15 // CleanupSuite adds the ability to add cleanup functions that are called 16 // during either test tear down or suite tear down depending on the method 17 // called. 18 type CleanupSuite struct { 19 testStack cleanupStack 20 suiteStack cleanupStack 21 } 22 23 func (s *CleanupSuite) SetUpSuite(c *gc.C) { 24 s.suiteStack = nil 25 } 26 27 func (s *CleanupSuite) TearDownSuite(c *gc.C) { 28 s.callStack(c, s.suiteStack) 29 } 30 31 func (s *CleanupSuite) SetUpTest(c *gc.C) { 32 s.testStack = nil 33 } 34 35 func (s *CleanupSuite) TearDownTest(c *gc.C) { 36 s.callStack(c, s.testStack) 37 } 38 39 func (s *CleanupSuite) callStack(c *gc.C, stack cleanupStack) { 40 for i := len(stack) - 1; i >= 0; i-- { 41 stack[i](c) 42 } 43 } 44 45 // AddCleanup pushes the cleanup function onto the stack of functions to be 46 // called during TearDownTest. 47 func (s *CleanupSuite) AddCleanup(cleanup CleanupFunc) { 48 s.testStack = append(s.testStack, cleanup) 49 } 50 51 // AddSuiteCleanup pushes the cleanup function onto the stack of functions to 52 // be called during TearDownSuite. 53 func (s *CleanupSuite) AddSuiteCleanup(cleanup CleanupFunc) { 54 s.suiteStack = append(s.suiteStack, cleanup) 55 } 56 57 // PatchEnvironment sets the environment variable 'name' the the value passed 58 // in. The old value is saved and returned to the original value at test tear 59 // down time using a cleanup function. 60 func (s *CleanupSuite) PatchEnvironment(name, value string) { 61 restore := PatchEnvironment(name, value) 62 s.AddCleanup(func(*gc.C) { restore() }) 63 } 64 65 // PatchEnvPathPrepend prepends the given path to the environment $PATH and restores the 66 // original path on test teardown. 67 func (s *CleanupSuite) PatchEnvPathPrepend(dir string) { 68 restore := PatchEnvPathPrepend(dir) 69 s.AddCleanup(func(*gc.C) { restore() }) 70 } 71 72 // PatchValue sets the 'dest' variable the the value passed in. The old value 73 // is saved and returned to the original value at test tear down time using a 74 // cleanup function. The value must be assignable to the element type of the 75 // destination. 76 func (s *CleanupSuite) PatchValue(dest, value interface{}) { 77 restore := PatchValue(dest, value) 78 s.AddCleanup(func(*gc.C) { restore() }) 79 } 80 81 // HookCommandOutput calls the package function of the same name to mock out 82 // the result of a particular comand execution, and will call the restore 83 // function on test teardown. 84 func (s *CleanupSuite) HookCommandOutput( 85 outputFunc *func(cmd *exec.Cmd) ([]byte, error), 86 output []byte, 87 err error, 88 ) <-chan *exec.Cmd { 89 result, restore := HookCommandOutput(outputFunc, output, err) 90 s.AddCleanup(func(*gc.C) { restore() }) 91 return result 92 }