launchpad.net/~rogpeppe/juju-core/500-errgo-fix@v0.0.0-20140213181702-000000002356/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  // PatchValue sets the 'dest' variable the the value passed in. The old value
    66  // is saved and returned to the original value at test tear down time using a
    67  // cleanup function. The value must be assignable to the element type of the
    68  // destination.
    69  func (s *CleanupSuite) PatchValue(dest, value interface{}) {
    70  	restore := PatchValue(dest, value)
    71  	s.AddCleanup(func(*gc.C) { restore() })
    72  }
    73  
    74  // HookCommandOutput calls the package function of the same name to mock out
    75  // the result of a particular comand execution, and will call the restore
    76  // function on test teardown.
    77  func (s *CleanupSuite) HookCommandOutput(
    78  	outputFunc *func(cmd *exec.Cmd) ([]byte, error),
    79  	output []byte,
    80  	err error,
    81  ) <-chan *exec.Cmd {
    82  	result, restore := HookCommandOutput(outputFunc, output, err)
    83  	s.AddCleanup(func(*gc.C) { restore() })
    84  	return result
    85  }