From 09379d03814c30f087e9da59b533343a9ef3c741 Mon Sep 17 00:00:00 2001
From: Gijs van Dam <gijs@lightning.engineering>
Date: Wed, 17 Dec 2025 17:09:00 +0100
Subject: [PATCH] wtclient: fix race conditions locking tests

This commit fixes two flaky test scenarios:

1. testRemoveLockedAddr: Add synchronization to wait for the dial to
   start before asserting that the address is locked. Previously, the
   test could race and check the lock state before session negotiation
   began.

2. testTowerSwitch: Use wait.Predicate for RemoveTower since the
   address may still be locked by an active session, causing
   intermittent failures.
---
 watchtower/wtclient/client_test.go | 33 ++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/watchtower/wtclient/client_test.go b/watchtower/wtclient/client_test.go
index e842876b6..1b50600ae 100644
--- a/watchtower/wtclient/client_test.go
+++ b/watchtower/wtclient/client_test.go
@@ -1811,8 +1811,10 @@ var clientTests = []clientTest{
 			require.NoError(h.t, err)
 
 			cancel := make(chan struct{})
+			dialStarted := make(chan struct{})
 			h.net.registerConnCallback(
 				h.server.addr, func(peer wtserver.Peer) {
+					close(dialStarted)
 					select {
 					case <-h.quit:
 					case <-cancel:
@@ -1847,6 +1849,16 @@ var clientTests = []clientTest{
 			err = h.clientMgr.AddTower(towerAddr)
 			require.NoError(h.t, err)
 
+			// Wait for the dial to start so that we know the
+			// session negotiation has begun and the address is
+			// locked.
+			select {
+			case <-dialStarted:
+
+			case <-time.After(waitTime):
+				h.t.Fatal("timeout waiting for dial to start")
+			}
+
 			// Assert that if the client attempts to remove the
 			// tower's first address, then it will error due to
 			// address currently being locked for session
@@ -2354,10 +2366,23 @@ var clientTests = []clientTest{
 			}, waitTime)
 			require.NoError(h.t, err)
 
-			// Now remove the tower.
-			err = h.clientMgr.RemoveTower(
-				h.server.addr.IdentityKey, nil,
-			)
+			// Now remove the tower. We use wait.Predicate here
+			// because the address may still be locked by an active
+			// session.
+			err = wait.Predicate(func() bool {
+				err := h.clientMgr.RemoveTower(
+					h.server.addr.IdentityKey, nil,
+				)
+				if err != nil {
+					require.ErrorIs(
+						h.t, err, wtclient.ErrAddrInUse,
+					)
+
+					return false
+				}
+
+				return true
+			}, waitTime)
 			require.NoError(h.t, err)
 
 			// Add a new tower.