Add NewT which automatically cleans up using t.Cleanup

gostub.go

@@ -34,13 +34,23 @@ type Stubs struct {
 }
 
 // New returns Stubs that can be used to stub out variables.
+//
+// Prefer `NewT` when using in tests, as it automatically resets.
 func New() *Stubs {
 	return &Stubs{
 		stubs:   make(map[reflect.Value]reflect.Value),
 		origEnv: make(map[string]envVal),
 	}
 }
 
+// NewT returns stubs that can be used to stub out variable, and automatically
+// resets stubs at the end of the test using `t.Cleanup`.
+func NewT(t TestingT) *Stubs {
+	stubs := New()
+	t.Cleanup(stubs.Reset)
+	return stubs
+}
+
 // 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.

gostub_test.go

@@ -23,6 +23,15 @@ func TestStub(t *testing.T) {
 	expectVal(t, v1, 100)
 }
 
+func TestNewT(t *testing.T) {
+	resetVars()
+	t.Run("inner", func(t *testing.T) {
+		// No need for reset.
+		NewT(t).Stub(&v1, 1)
+	})
+	expectVal(t, 100, v1)
+}
+
 func TestValue(t *testing.T) {
 	resetVars()