WIP

doc.go

@@ -5,78 +5,86 @@ original value once the test has been run.
 This can be used to stub static variables as well as static functions. To
 stub a static variable, use the Stub function:
 
-  var configFile = "config.json"
+	var configFile = "config.json"
 
-  func GetConfig() ([]byte, error) {
-    return ioutil.ReadFile(configFile)
-  }
+	func GetConfig() ([]byte, error) {
+	  return ioutil.ReadFile(configFile)
+	}
 
-  // Test code
-  stubs := gostub.Stub(&configFile, "/tmp/test.config")
+	// Test code
+	stubs := gostub.Stub(t, &configFile, "/tmp/test.config")
 
-  data, err := GetConfig()
-  // data will now return contents of the /tmp/test.config file
+	data, err := GetConfig()
+	// data will now return contents of the /tmp/test.config file
 
 gostub can also stub static functions in a test by using a variable
 to reference the static function, and using that local variable to call
 the static function:
-  var timeNow = time.Now
 
-  func GetDate() int {
-  	return timeNow().Day()
-  }
+	var timeNow = time.Now
+
+	func GetDate() int {
+		return timeNow().Day()
+	}
 
 You can test this by using gostub to stub the timeNow variable:
-  stubs := gostub.Stub(&timeNow, func() time.Time {
-    return time.Date(2015, 6, 1, 0, 0, 0, 0, time.UTC)
-  })
-  defer stubs.Reset()
 
-  // Test can check that GetDate returns 6
+	stubs := gostub.Stub(&timeNow, func() time.Time {
+	  return time.Date(2015, 6, 1, 0, 0, 0, 0, time.UTC)
+	})
+	defer stubs.Reset()
+
+	// Test can check that GetDate returns 6
 
 If you are stubbing a function to return a constant value like in
 the above test, you can use StubFunc instead:
-  stubs := gostub.StubFunc(&timeNow, time.Date(2015, 6, 1, 0, 0, 0, 0, time.UTC))
-  defer stubs.Reset()
+
+	stubs := gostub.StubFunc(&timeNow, time.Date(2015, 6, 1, 0, 0, 0, 0, time.UTC))
+	defer stubs.Reset()
 
 StubFunc can also be used to stub functions that return multiple values:
-  var osHostname = osHostname
-  // [...] production code using osHostname to call it.
 
-  // Test code:
-  stubs := gostub.StubFunc(&osHostname, "fakehost", nil)
-  defer stubs.Reset()
+	var osHostname = osHostname
+	// [...] production code using osHostname to call it.
+
+	// Test code:
+	stubs := gostub.StubFunc(&osHostname, "fakehost", nil)
+	defer stubs.Reset()
 
 StubEnv can be used to setup environment variables for tests, and the environment
 values are reset to their original values upon Reset:
 
-  stubs := gostub.New()
-  stubs.SetEnv("GOSTUB_VAR", "test_value")
-  defer stubs.Reset()
+	stubs := gostub.New()
+	stubs.SetEnv("GOSTUB_VAR", "test_value")
+	defer stubs.Reset()
 
 The Reset method should be deferred to run at the end of the test to reset
 all stubbed variables back to their original values.
 
 You can set up multiple stubs by calling Stub again:
-  stubs := gostub.Stub(&v1, 1)
-  stubs.Stub(&v2, 2)
-  defer stubs.Reset()
+
+	stubs := gostub.Stub(&v1, 1)
+	stubs.Stub(&v2, 2)
+	defer stubs.Reset()
 
 For simple cases where you are only setting up simple stubs, you can condense
 the setup and cleanup into a single line:
-  defer gostub.Stub(&v1, 1).Stub(&v2, 2).Reset()
+
+	defer gostub.Stub(&v1, 1).Stub(&v2, 2).Reset()
+
 This sets up the stubs and then defers the Reset call.
 
 You should keep the return argument from the Stub call if you need to change
 stubs or add more stubs during test execution:
-  stubs := gostub.Stub(&v1, 1)
-  defer stubs.Reset()
 
-  // Do some testing
-  stubs.Stub(&v1, 5)
+	stubs := gostub.Stub(&v1, 1)
+	defer stubs.Reset()
+
+	// Do some testing
+	stubs.Stub(&v1, 5)
 
-  // More testing
-  stubs.Stub(&b2, 6)
+	// More testing
+	stubs.Stub(&b2, 6)
 
 The Stub call must be passed a pointer to the variable that should be stubbed,
 and a value which can be assigned to the variable.

gostub.go

@@ -1,84 +1,26 @@
 package gostub
 
 import (
-	"fmt"
 	"reflect"
+	"testing"
 )
 
-// Stub replaces the value stored at varToStub with stubVal.
-// varToStub must be a pointer to the variable. stubVal should have a type
-// that is assignable to the variable.
-func Stub(varToStub interface{}, stubVal interface{}) *Stubs {
-	return New().Stub(varToStub, stubVal)
-}
-
-// StubFunc replaces a function variable with a function that returns stubVal.
-// funcVarToStub must be a pointer to a function variable. If the function
-// returns multiple values, then multiple values should be passed to stubFunc.
-// The values must match be assignable to the return values' types.
-func StubFunc(funcVarToStub interface{}, stubVal ...interface{}) *Stubs {
-	return New().StubFunc(funcVarToStub, stubVal...)
-}
-
-type envVal struct {
-	val string
-	ok  bool
-}
-
-// Stubs represents a set of stubbed variables that can be reset.
-type Stubs struct {
-	// stubs is a map from the variable pointer (being stubbed) to the original value.
-	stubs   map[reflect.Value]reflect.Value
-	origEnv map[string]envVal
-}
-
-// New returns Stubs that can be used to stub out variables.
-func New() *Stubs {
-	return &Stubs{
-		stubs:   make(map[reflect.Value]reflect.Value),
-		origEnv: make(map[string]envVal),
-	}
-}
-
-// Stub replaces the value stored at varToStub with stubVal.
-// varToStub must be a pointer to the variable. stubVal should have a type
-// that is assignable to the variable.
-func (s *Stubs) Stub(varToStub interface{}, stubVal interface{}) *Stubs {
-	v := reflect.ValueOf(varToStub)
-	stub := reflect.ValueOf(stubVal)
-
-	// Ensure varToStub is a pointer to the variable.
-	if v.Type().Kind() != reflect.Ptr {
-		panic("variable to stub is expected to be a pointer")
-	}
-
-	if _, ok := s.stubs[v]; !ok {
-		// Store the original value if this is the first time varPtr is being stubbed.
-		s.stubs[v] = reflect.ValueOf(v.Elem().Interface())
+// Stub replaces the value stored at varPtrToStub with stubVal.
+// varToStub must be a pointer to the variable.
+// The variable will be reset to the original value at the end of the test
+// but it can be reset earlier using the returned function.
+func Stub[T any](t testing.TB, varPtrToStub *T, stubVal T) (reset func()) {
+	orig := *varPtrToStub
+	*varPtrToStub = stubVal
+	cleanup := func() {
+		*varPtrToStub = orig
 	}
-
-	// *varToStub = stubVal
-	v.Elem().Set(stub)
-	return s
+	t.Cleanup(cleanup)
+	return cleanup
 }
 
-// StubFunc replaces a function variable with a function that returns stubVal.
-// funcVarToStub must be a pointer to a function variable. If the function
-// returns multiple values, then multiple values should be passed to stubFunc.
-// The values must match be assignable to the return values' types.
-func (s *Stubs) StubFunc(funcVarToStub interface{}, stubVal ...interface{}) *Stubs {
-	funcPtrType := reflect.TypeOf(funcVarToStub)
-	if funcPtrType.Kind() != reflect.Ptr ||
-		funcPtrType.Elem().Kind() != reflect.Func {
-		panic("func variable to stub must be a pointer to a function")
-	}
-	funcType := funcPtrType.Elem()
-	if funcType.NumOut() != len(stubVal) {
-		panic(fmt.Sprintf("func type has %v return values, but only %v stub values provided",
-			funcType.NumOut(), len(stubVal)))
-	}
-
-	return s.Stub(funcVarToStub, FuncReturning(funcPtrType.Elem(), stubVal...).Interface())
+type Func[T any] interface {
+	func(any) T | func(any, any) T
 }
 
 // FuncReturning creates a new function with type funcType that returns results.
@@ -106,14 +48,6 @@ func FuncReturning(funcType reflect.Type, results ...interface{}) reflect.Value
 	})
 }
 
-// Reset resets all stubbed variables back to their original values.
-func (s *Stubs) Reset() {
-	for v, originalVal := range s.stubs {
-		v.Elem().Set(originalVal)
-	}
-	s.resetEnv()
-}
-
 // ResetSingle resets a single stubbed variable back to its original value.
 func (s *Stubs) ResetSingle(varToStub interface{}) {
 	v := reflect.ValueOf(varToStub)