tstest: add method to Replace values for tests

We have many function pointers that we replace for the duration of test and
restore it on test completion, add method to do that.

Signed-off-by: Maisem Ali <maisem@tailscale.com>

tstest/tstest.go

@@ -6,12 +6,29 @@ package tstest
 
 import (
 	"context"
+	"testing"
 	"time"
 
 	"tailscale.com/logtail/backoff"
 	"tailscale.com/types/logger"
 )
 
+// Replace replaces the value of target with val.
+// The old value is restored when the test ends.
+func Replace[T any](t *testing.T, target *T, val T) {
+	t.Helper()
+	if target == nil {
+		t.Fatalf("Replace: nil pointer")
+	}
+	old := *target
+	t.Cleanup(func() {
+		*target = old
+	})
+
+	*target = val
+	return
+}
+
 // WaitFor retries try for up to maxWait.
 // It returns nil once try returns nil the first time.
 // If maxWait passes without success, it returns try's last error.

tstest/tstest_test.go

@@ -0,0 +1,24 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package tstest
+
+import "testing"
+
+func TestReplace(t *testing.T) {
+	before := "before"
+	done := false
+	t.Run("replace", func(t *testing.T) {
+		Replace(t, &before, "after")
+		if before != "after" {
+			t.Errorf("before = %q; want %q", before, "after")
+		}
+		done = true
+	})
+	if !done {
+		t.Fatal("subtest didn't run")
+	}
+	if before != "before" {
+		t.Errorf("before = %q; want %q", before, "before")
+	}
+}