From 43375c6efbfbcb2dbbe58107bc0a2cabec173668 Mon Sep 17 00:00:00 2001
From: Nick Khyl <nickk@tailscale.com>
Date: Tue, 23 Jul 2024 19:34:03 -0500
Subject: [PATCH] types/lazy: re-init SyncValue during test cleanup if it
 wasn't set before SetForTest

Updates #12687

Signed-off-by: Nick Khyl <nickk@tailscale.com>
---
 types/lazy/lazy.go      | 10 +++++++---
 types/lazy/sync_test.go | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/types/lazy/lazy.go b/types/lazy/lazy.go
index 8bd55bdf6..43325512d 100644
--- a/types/lazy/lazy.go
+++ b/types/lazy/lazy.go
@@ -170,8 +170,8 @@ type TB interface {
 func (z *SyncValue[T]) SetForTest(tb TB, val T, err error) {
 	tb.Helper()
 
-	z.once.Do(func() {})
 	oldErr, oldVal := z.err.Load(), z.v
+	z.once.Do(func() {})
 
 	z.v = val
 	if err != nil {
@@ -181,7 +181,11 @@ func (z *SyncValue[T]) SetForTest(tb TB, val T, err error) {
 	}
 
 	tb.Cleanup(func() {
-		z.v = oldVal
-		z.err.Store(oldErr)
+		if oldErr == nil {
+			*z = SyncValue[T]{}
+		} else {
+			z.v = oldVal
+			z.err.Store(oldErr)
+		}
 	})
 }
diff --git a/types/lazy/sync_test.go b/types/lazy/sync_test.go
index 8fdf9e76f..5578eee0c 100644
--- a/types/lazy/sync_test.go
+++ b/types/lazy/sync_test.go
@@ -290,6 +290,22 @@ func TestSyncValueSetForTest(t *testing.T) {
 						t.Fatalf("CleanupValue: got %v; want %v", gotCleanupValue, wantCleanupValue)
 					}
 				})
+			} else {
+				// Verify that if v wasn't set prior to SetForTest, it's
+				// reverted to a valid unset state during the test cleanup.
+				t.Cleanup(func() {
+					if _, _, ok := v.PeekErr(); ok {
+						t.Fatal("SyncValue is set after cleanup")
+					}
+					wantCleanupValue, wantCleanupErr := 42, errors.New("ka-boom")
+					gotCleanupValue, gotCleanupErr := v.GetErr(func() (int, error) { return wantCleanupValue, wantCleanupErr })
+					if gotCleanupErr != wantCleanupErr {
+						t.Fatalf("CleanupErr: got %v; want %v", gotCleanupErr, wantCleanupErr)
+					}
+					if gotCleanupValue != wantCleanupValue {
+						t.Fatalf("CleanupValue: got %v; want %v", gotCleanupValue, wantCleanupValue)
+					}
+				})
 			}
 
 			// Set the test value and/or error.