From 1b25d8a1223734e63b48638e34cc06614d34298e Mon Sep 17 00:00:00 2001
From: Boris Nagaev <bnagaev@gmail.com>
Date: Fri, 15 May 2026 18:29:11 -0500
Subject: [PATCH] sweepbatcher: assert close-during-add run exit

TestSweepBatcherCloseDuringAdding previously started Batcher.Run in a
detached goroutine and called test assertions from that goroutine. The
test only waited for the add/cancel workers, so a Run-side shutdown
error could be missed or reported unreliably.

Route the Run result through a channel and wait for it in the main test
goroutine. While waiting, keep draining spend registrations so shutdown
cannot deadlock on mock notifier traffic.

This makes the existing shutdown-race test cover both sides of the race:
AddSweep callers may exit with cancellation, and Batcher.Run must also
terminate with an expected shutdown error.
---
 sweepbatcher/sweep_batcher_test.go | 37 +++++++++++++++++-------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/sweepbatcher/sweep_batcher_test.go b/sweepbatcher/sweep_batcher_test.go
index 648dc4c6..bfca8bb5 100644
--- a/sweepbatcher/sweep_batcher_test.go
+++ b/sweepbatcher/sweep_batcher_test.go
@@ -4208,9 +4208,9 @@ func testSweepBatcherCloseDuringAdding(t *testing.T, store testStore,
 	batcher := NewBatcher(lnd.WalletKit, lnd.ChainNotifier, lnd.Signer,
 		testMuSig2SignSweep, testVerifySchnorrSig, lnd.ChainParams,
 		batcherStore, sweepStore)
+	runErrChan := make(chan error, 1)
 	go func() {
-		err := batcher.Run(ctx)
-		checkBatcherError(t, err)
+		runErrChan <- batcher.Run(ctx)
 	}()
 
 	// Add many swaps.
@@ -4276,23 +4276,28 @@ func testSweepBatcherCloseDuringAdding(t *testing.T, store testStore,
 	})
 
 	// We don't know how many spend notification registrations will be
-	// issued, so accept them while waiting for two goroutines to stop.
-	quit := make(chan struct{})
-	registrationChan := make(chan struct{})
+	// issued, so accept them while waiting for all goroutines to stop.
+	addDone := make(chan struct{})
 	go func() {
-		defer close(registrationChan)
-		for {
-			select {
-			case <-lnd.RegisterSpendChannel:
-			case <-quit:
-				return
-			}
-		}
+		defer close(addDone)
+		wg.Wait()
 	}()
 
-	wg.Wait()
-	close(quit)
-	<-registrationChan
+	for addDone != nil || runErrChan != nil {
+		select {
+		case <-lnd.RegisterSpendChannel:
+
+		case <-addDone:
+			addDone = nil
+
+		case err := <-runErrChan:
+			checkBatcherError(t, err)
+			runErrChan = nil
+
+		case <-time.After(test.Timeout):
+			t.Fatalf("expected batcher close during adding to finish")
+		}
+	}
 }
 
 // testSweepBatcherHandleSweepRace reproduces a race between AddSweep and the