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  }