ipn/ipnlocal, util/syspolicy: convert last RegisterWellKnownSettingsForTest caller, remove

Updates #16998

Change-Id: I735d75129a97a929092e9075107e41cdade18944
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

5 files changed, 129 insertions, 36 deletions

diff --git a/ipn/ipnlocal/local_test.go b/ipn/ipnlocal/local_test.go
index 4debcdd8d..7d1c452f3 100644
--- a/ipn/ipnlocal/local_test.go
+++ b/ipn/ipnlocal/local_test.go
@@ -65,7 +65,6 @@ import (
 	"tailscale.com/util/syspolicy"
 	"tailscale.com/util/syspolicy/pkey"
 	"tailscale.com/util/syspolicy/policytest"
-	"tailscale.com/util/syspolicy/setting"
 	"tailscale.com/util/syspolicy/source"
 	"tailscale.com/wgengine"
 	"tailscale.com/wgengine/filter"
@@ -6529,12 +6528,13 @@ func TestUpdatePrefsOnSysPolicyChange(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			syspolicy.RegisterWellKnownSettingsForTest(t)
-			store := source.NewTestStoreOf[string](t)
-			syspolicy.MustRegisterStoreForTest(t, "TestSource", setting.DeviceScope, store)
+			var polc policytest.Config
+			polc.EnableRegisterChangeCallback()
 
 			sys := tsd.NewSystem()
+			sys.PolicyClient.Set(polc)
 			lb := newLocalBackendWithSysAndTestControl(t, enableLogging, sys, func(tb testing.TB, opts controlclient.Options) controlclient.Client {
+				opts.PolicyClient = polc
 				return newClient(tb, opts)
 			})
 			if tt.initialPrefs != nil {
@@ -6551,7 +6551,11 @@ func TestUpdatePrefsOnSysPolicyChange(t *testing.T) {
 				nw.watch(0, nil, unexpectedPrefsChange)
 			}
 
-			store.SetStrings(tt.stringSettings...)
+			var batch policytest.Config
+			for _, ss := range tt.stringSettings {
+				batch.Set(ss.Key, ss.Value)
+			}
+			polc.SetMultiple(batch)
 
 			nw.check()
 		})
diff --git a/util/syspolicy/policy_keys.go b/util/syspolicy/policy_keys.go
index 1bbcfe6ca..ef2ac430d 100644
--- a/util/syspolicy/policy_keys.go
+++ b/util/syspolicy/policy_keys.go
@@ -76,13 +76,3 @@ func init() {
 		return nil
 	})
 }
-
-// RegisterWellKnownSettingsForTest registers all implicit setting definitions
-// for the duration of the test.
-func RegisterWellKnownSettingsForTest(tb testenv.TB) {
-	tb.Helper()
-	err := setting.SetDefinitionsForTest(tb, implicitDefinitions...)
-	if err != nil {
-		tb.Fatalf("Failed to register well-known settings: %v", err)
-	}
-}
diff --git a/util/syspolicy/policytest/policytest.go b/util/syspolicy/policytest/policytest.go
index 7ea0ad91f..e5c1c7856 100644
--- a/util/syspolicy/policytest/policytest.go
+++ b/util/syspolicy/policytest/policytest.go
@@ -6,8 +6,12 @@ package policytest
 
 import (
 	"fmt"
+	"maps"
+	"slices"
+	"sync"
 	"time"
 
+	"tailscale.com/util/set"
 	"tailscale.com/util/syspolicy/pkey"
 	"tailscale.com/util/syspolicy/policyclient"
 	"tailscale.com/util/syspolicy/ptype"
@@ -29,11 +33,85 @@ type Config map[pkey.Key]any
 
 var _ policyclient.Client = Config{}
 
+// Set sets key to value. The value should be of the correct type that it will
+// be read as later. For PreferenceOption and Visibility, you may also set them
+// to 'string' values and they'll be UnmarshalText'ed into their correct value
+// at Get time.
+//
+// As a special case, the value can also be of type error to make the accessors
+// return that error value.
 func (c *Config) Set(key pkey.Key, value any) {
 	if *c == nil {
 		*c = make(map[pkey.Key]any)
 	}
 	(*c)[key] = value
+
+	if w, ok := (*c)[watchersKey].(*watchers); ok && key != watchersKey {
+		w.mu.Lock()
+		vals := slices.Collect(maps.Values(w.s))
+		w.mu.Unlock()
+		for _, f := range vals {
+			f(policyChange(key))
+		}
+	}
+}
+
+// SetMultiple is a batch version of [Config.Set]. It copies the contents of o
+// into c and does at most one notification wake-up for the whole batch.
+func (c *Config) SetMultiple(o Config) {
+	if *c == nil {
+		*c = make(map[pkey.Key]any)
+	}
+
+	maps.Copy(*c, o)
+
+	if w, ok := (*c)[watchersKey].(*watchers); ok {
+		w.mu.Lock()
+		vals := slices.Collect(maps.Values(w.s))
+		w.mu.Unlock()
+		for _, f := range vals {
+			f(policyChanges(o))
+		}
+	}
+}
+
+type policyChange pkey.Key
+
+func (pc policyChange) HasChanged(v pkey.Key) bool { return pkey.Key(pc) == v }
+func (pc policyChange) HasChangedAnyOf(keys ...pkey.Key) bool {
+	return slices.Contains(keys, pkey.Key(pc))
+}
+
+type policyChanges map[pkey.Key]any
+
+func (pc policyChanges) HasChanged(v pkey.Key) bool {
+	_, ok := pc[v]
+	return ok
+}
+func (pc policyChanges) HasChangedAnyOf(keys ...pkey.Key) bool {
+	for _, k := range keys {
+		if pc.HasChanged(k) {
+			return true
+		}
+	}
+	return false
+}
+
+const watchersKey = "_policytest_watchers"
+
+type watchers struct {
+	mu sync.Mutex
+	s  set.HandleSet[func(policyclient.PolicyChange)]
+}
+
+// EnableRegisterChangeCallback makes c support the RegisterChangeCallback
+// for testing. Without calling this, the RegisterChangeCallback does nothing.
+// For watchers to be notified, use the [Config.Set] method. Changing the map
+// directly obviously wouldn't work.
+func (c *Config) EnableRegisterChangeCallback() {
+	if _, ok := (*c)[watchersKey]; !ok {
+		c.Set(watchersKey, new(watchers))
+	}
 }
 
 func (c Config) GetStringArray(key pkey.Key, defaultVal []string) ([]string, error) {
@@ -153,8 +231,19 @@ func (c Config) HasAnyOf(keys ...pkey.Key) (bool, error) {
 	return false, nil
 }
 
-func (sp Config) RegisterChangeCallback(callback func(policyclient.PolicyChange)) (func(), error) {
-	return func() {}, nil
+func (c Config) RegisterChangeCallback(callback func(policyclient.PolicyChange)) (func(), error) {
+	w, ok := c[watchersKey].(*watchers)
+	if !ok {
+		return func() {}, nil
+	}
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	h := w.s.Add(callback)
+	return func() {
+		w.mu.Lock()
+		defer w.mu.Unlock()
+		delete(w.s, h)
+	}, nil
 }
 
 func (sp Config) SetDebugLoggingEnabled(enabled bool) {}
diff --git a/util/syspolicy/syspolicy.go b/util/syspolicy/syspolicy.go
index 2367e21eb..48e430b67 100644
--- a/util/syspolicy/syspolicy.go
+++ b/util/syspolicy/syspolicy.go
@@ -19,7 +19,6 @@ import (
 	"tailscale.com/util/syspolicy/rsop"
 	"tailscale.com/util/syspolicy/setting"
 	"tailscale.com/util/syspolicy/source"
-	"tailscale.com/util/testenv"
 )
 
 var (
@@ -45,16 +44,6 @@ func RegisterStore(name string, scope setting.PolicyScope, store source.Store) (
 	return rsop.RegisterStore(name, scope, store)
 }
 
-// MustRegisterStoreForTest is like [rsop.RegisterStoreForTest], but it fails the test if the store could not be registered.
-func MustRegisterStoreForTest(tb testenv.TB, name string, scope setting.PolicyScope, store source.Store) *rsop.StoreRegistration {
-	tb.Helper()
-	reg, err := rsop.RegisterStoreForTest(tb, name, scope, store)
-	if err != nil {
-		tb.Fatalf("Failed to register policy store %q as a %v policy source: %v", name, scope, err)
-	}
-	return reg
-}
-
 // hasAnyOf returns whether at least one of the specified policy settings is configured,
 // or an error if no keys are provided or the check fails.
 func hasAnyOf(keys ...pkey.Key) (bool, error) {
diff --git a/util/syspolicy/syspolicy_test.go b/util/syspolicy/syspolicy_test.go
index 0ee62efb1..10f8da486 100644
--- a/util/syspolicy/syspolicy_test.go
+++ b/util/syspolicy/syspolicy_test.go
@@ -14,6 +14,7 @@ import (
 	"tailscale.com/util/syspolicy/internal/metrics"
 	"tailscale.com/util/syspolicy/pkey"
 	"tailscale.com/util/syspolicy/ptype"
+	"tailscale.com/util/syspolicy/rsop"
 	"tailscale.com/util/syspolicy/setting"
 	"tailscale.com/util/syspolicy/source"
 	"tailscale.com/util/testenv"
@@ -21,6 +22,16 @@ import (
 
 var someOtherError = errors.New("error other than not found")
 
+// registerWellKnownSettingsForTest registers all implicit setting definitions
+// for the duration of the test.
+func registerWellKnownSettingsForTest(tb testenv.TB) {
+	tb.Helper()
+	err := setting.SetDefinitionsForTest(tb, implicitDefinitions...)
+	if err != nil {
+		tb.Fatalf("Failed to register well-known settings: %v", err)
+	}
+}
+
 func TestGetString(t *testing.T) {
 	tests := []struct {
 		name         string
@@ -68,7 +79,7 @@ func TestGetString(t *testing.T) {
 		},
 	}
 
-	RegisterWellKnownSettingsForTest(t)
+	registerWellKnownSettingsForTest(t)
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -210,7 +221,7 @@ func TestGetBoolean(t *testing.T) {
 		},
 	}
 
-	RegisterWellKnownSettingsForTest(t)
+	registerWellKnownSettingsForTest(t)
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -303,7 +314,7 @@ func TestGetPreferenceOption(t *testing.T) {
 		},
 	}
 
-	RegisterWellKnownSettingsForTest(t)
+	registerWellKnownSettingsForTest(t)
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -388,7 +399,7 @@ func TestGetVisibility(t *testing.T) {
 		},
 	}
 
-	RegisterWellKnownSettingsForTest(t)
+	registerWellKnownSettingsForTest(t)
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -484,7 +495,7 @@ func TestGetDuration(t *testing.T) {
 		},
 	}
 
-	RegisterWellKnownSettingsForTest(t)
+	registerWellKnownSettingsForTest(t)
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -565,7 +576,7 @@ func TestGetStringArray(t *testing.T) {
 		},
 	}
 
-	RegisterWellKnownSettingsForTest(t)
+	registerWellKnownSettingsForTest(t)
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -599,14 +610,24 @@ func TestGetStringArray(t *testing.T) {
 	}
 }
 
+// mustRegisterStoreForTest is like [rsop.RegisterStoreForTest], but it fails the test if the store could not be registered.
+func mustRegisterStoreForTest(tb testenv.TB, name string, scope setting.PolicyScope, store source.Store) *rsop.StoreRegistration {
+	tb.Helper()
+	reg, err := rsop.RegisterStoreForTest(tb, name, scope, store)
+	if err != nil {
+		tb.Fatalf("Failed to register policy store %q as a %v policy source: %v", name, scope, err)
+	}
+	return reg
+}
+
 func registerSingleSettingStoreForTest[T source.TestValueType](tb testenv.TB, s source.TestSetting[T]) {
 	policyStore := source.NewTestStoreOf(tb, s)
-	MustRegisterStoreForTest(tb, "TestStore", setting.DeviceScope, policyStore)
+	mustRegisterStoreForTest(tb, "TestStore", setting.DeviceScope, policyStore)
 }
 
 func BenchmarkGetString(b *testing.B) {
 	loggerx.SetForTest(b, logger.Discard, logger.Discard)
-	RegisterWellKnownSettingsForTest(b)
+	registerWellKnownSettingsForTest(b)
 
 	wantControlURL := "https://login.tailscale.com"
 	registerSingleSettingStoreForTest(b, source.TestSettingOf(pkey.ControlURL, wantControlURL))