From bb62eed5dbaf9c014a44c5cc71afbebacbb1ab3b Mon Sep 17 00:00:00 2001
From: David Reiss <david@temporal.io>
Date: Tue, 24 Jan 2023 18:51:55 -0800
Subject: [PATCH] Add RefreshPerNSWorkerManager and RemoveOverride for tests

---
 service/worker/pernamespaceworker.go | 14 +++++++++-----
 service/worker/service.go            |  5 +++++
 tests/dynamicconfig.go               |  6 ++++++
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/service/worker/pernamespaceworker.go b/service/worker/pernamespaceworker.go
index a8a9611cb5d..212c60a42a0 100644
--- a/service/worker/pernamespaceworker.go
+++ b/service/worker/pernamespaceworker.go
@@ -194,13 +194,17 @@ func (wm *perNamespaceWorkerManager) namespaceCallback(ns *namespace.Namespace,
 	go wm.getWorkerByNamespace(ns).refreshWithNewNamespace(ns, deleted)
 }
 
+func (wm *perNamespaceWorkerManager) refreshAll() {
+	wm.lock.Lock()
+	defer wm.lock.Unlock()
+	for _, worker := range wm.workers {
+		go worker.refreshWithExistingNamespace()
+	}
+}
+
 func (wm *perNamespaceWorkerManager) membershipChangedListener() {
 	for range wm.membershipChangedCh {
-		wm.lock.Lock()
-		for _, worker := range wm.workers {
-			go worker.refreshWithExistingNamespace()
-		}
-		wm.lock.Unlock()
+		wm.refreshAll()
 	}
 }
 
diff --git a/service/worker/service.go b/service/worker/service.go
index a246d8d1f5a..79dac551038 100644
--- a/service/worker/service.go
+++ b/service/worker/service.go
@@ -564,3 +564,8 @@ func (s *Service) ensureSystemNamespaceExists(
 		)
 	}
 }
+
+// This is intended for use by integration tests only.
+func (s *Service) RefreshPerNSWorkerManager() {
+	s.perNamespaceWorkerManager.refreshAll()
+}
diff --git a/tests/dynamicconfig.go b/tests/dynamicconfig.go
index b5fca93bd33..394a544ccde 100644
--- a/tests/dynamicconfig.go
+++ b/tests/dynamicconfig.go
@@ -81,6 +81,12 @@ func (d *dcClient) OverrideValue(name dynamicconfig.Key, value any) {
 	d.overrides[name] = value
 }
 
+func (d *dcClient) RemoveOverride(name dynamicconfig.Key) {
+	d.Lock()
+	defer d.Unlock()
+	delete(d.overrides, name)
+}
+
 // newTestDCClient - returns a dynamic config client for integration testing
 func newTestDCClient(fallback dynamicconfig.Client) *dcClient {
 	return &dcClient{