From 81b9394525e3ec4398f62433e5e58a04b9af9dd7 Mon Sep 17 00:00:00 2001
From: Prashant V <prashant@prashantv.com>
Date: Tue, 8 Oct 2024 22:36:08 -0700
Subject: [PATCH] Use generics to stub a single value with auto-reset

---
 gostub.go      |  1 +
 gostub_test.go | 30 +++++++++++++++++++++++++++---
 value.go       | 20 ++++++++++++++++++++
 3 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100644 value.go

diff --git a/gostub.go b/gostub.go
index ce9fd84..e67b10a 100644
--- a/gostub.go
+++ b/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)
 }
diff --git a/gostub_test.go b/gostub_test.go
index 35d1476..fded595 100644
--- a/gostub_test.go
+++ b/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()
 
diff --git a/value.go b/value.go
new file mode 100644
index 0000000..28a6348
--- /dev/null
+++ b/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
+}