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.

Author: gijswijs

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,15 @@ 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 +2365,16 @@ 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,
+				)
+
+				return err == nil
+			}, waitTime)
 			require.NoError(h.t, err)
 
 			// Add a new tower.