Use generics to stub a single value with auto-reset (#21)

gostub.go

@@ -8,6 +8,7 @@ import (
 // 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.
+// When stubbing a single value, prefer `Value`.
 func Stub(varToStub interface{}, stubVal interface{}) *Stubs {
 	return New().Stub(varToStub, stubVal)
 }

gostub_test.go

@@ -18,14 +18,38 @@ func TestStub(t *testing.T) {
 
 	stubs := Stub(&v1, 1)
 
-	if v1 != 1 {
-		t.Errorf("expected")
-	}
 	expectVal(t, v1, 1)
 	stubs.Reset()
 	expectVal(t, v1, 100)
 }
 
+func TestValue(t *testing.T) {
+	resetVars()
+
+	t.Run("test", func(t *testing.T) {
+		reset1 := Value(t, &v1, 1)
+		reset2 := Value(t, &v2, 2)
+		expectVal(t, v1, 1)
+		expectVal(t, v2, 2)
+		reset1()
+		expectVal(t, v1, 100)
+		expectVal(t, v2, 2)
+		reset2()
+		expectVal(t, v1, 100)
+		expectVal(t, v2, 200)
+
+		Value(t, &v1, 0)
+		Value(t, &v2, 0)
+		Value(t, &v3, 0)
+	})
+
+	t.Run("verify Cleanup", func(t *testing.T) {
+		expectVal(t, v1, 100)
+		expectVal(t, v2, 200)
+		expectVal(t, v3, 300)
+	})
+}
+
 func TestRestub(t *testing.T) {
 	resetVars()
 

value.go

@@ -0,0 +1,20 @@
+package gostub
+
+// TestingT is a subset of the testing.TB interface used by gostub.
+type TestingT interface {
+	Cleanup(func())
+}
+
+// Value replaces the value at varPtr with stubVal.
+// The original value is reset at the end of the test via t.Cleanup
+// or can be reset using the returned function.
+func Value[T any](t TestingT, varPtr *T, stubVal T) (reset func()) {
+	orig := *varPtr
+	*varPtr = stubVal
+
+	reset = func() {
+		*varPtr = orig
+	}
+	t.Cleanup(reset)
+	return reset
+}