From 440779fc8cd876b1381e391f914f65dcb7609341 Mon Sep 17 00:00:00 2001
From: Ivan Belyakov <ivan.belyakov.job@gmail.com>
Date: Sun, 3 Mar 2024 08:59:21 +0100
Subject: [PATCH] fix: flaky test TestController_watchAccountChanges by making
 sure to check the database transfer after removal event is processed.

---
 services/wallet/transfer/controller.go      |  2 +-
 services/wallet/transfer/controller_test.go | 63 +++++++++++----------
 2 files changed, 33 insertions(+), 32 deletions(-)

diff --git a/services/wallet/transfer/controller.go b/services/wallet/transfer/controller.go
index 4ac717d5d..14d1051dc 100644
--- a/services/wallet/transfer/controller.go
+++ b/services/wallet/transfer/controller.go
@@ -150,8 +150,8 @@ func (c *Controller) startAccountWatcher(chainIDs []uint64) {
 		c.accWatcher = accountsevent.NewWatcher(c.accountsDB, c.accountFeed, func(changedAddresses []common.Address, eventType accountsevent.EventType, currentAddresses []common.Address) {
 			c.onAccountsChanged(changedAddresses, eventType, currentAddresses, chainIDs)
 		})
-		c.accWatcher.Start()
 	}
+	c.accWatcher.Start()
 }
 
 func (c *Controller) onAccountsChanged(changedAddresses []common.Address, eventType accountsevent.EventType, currentAddresses []common.Address, chainIDs []uint64) {
diff --git a/services/wallet/transfer/controller_test.go b/services/wallet/transfer/controller_test.go
index 6a0798879..1eea23ce8 100644
--- a/services/wallet/transfer/controller_test.go
+++ b/services/wallet/transfer/controller_test.go
@@ -80,12 +80,40 @@ func TestController_watchAccountsChanges(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, ranges)
 
-	ch := make(chan accountsevent.Event)
-	// Subscribe for account changes
-	accountFeed.Subscribe(ch)
+	c.accWatcher = accountsevent.NewWatcher(c.accountsDB, c.accountFeed, func(changedAddresses []common.Address, eventType accountsevent.EventType, currentAddresses []common.Address) {
+		c.onAccountsChanged(changedAddresses, eventType, currentAddresses, []uint64{chainID})
+
+		// Quit channel event handler before  destroying the channel
+		go func() {
+			time.Sleep(1 * time.Millisecond)
+			// Wait for DB to be cleaned up
+			c.accWatcher.Stop()
+
+			// Check that transfers, blocks and block ranges were deleted
+			transfers, err := database.GetTransfersByAddress(chainID, address, big.NewInt(2), 1)
+			require.NoError(t, err)
+			require.Len(t, transfers, 0)
+
+			blocksDAO := &BlockDAO{walletDB}
+			block, err := blocksDAO.GetLastBlockByAddress(chainID, address, 1)
+			require.NoError(t, err)
+			require.Nil(t, block)
+
+			ranges, _, err = blockRangesDAO.getBlockRange(chainID, address)
+			require.NoError(t, err)
+			require.Nil(t, ranges.eth.FirstKnown)
+			require.Nil(t, ranges.eth.LastKnown)
+			require.Nil(t, ranges.eth.Start)
+			require.Nil(t, ranges.tokens.FirstKnown)
+			require.Nil(t, ranges.tokens.LastKnown)
+			require.Nil(t, ranges.tokens.Start)
+		}()
+	})
+	c.startAccountWatcher([]uint64{chainID})
 
 	// Watching accounts must start before sending event.
-	// To avoid running goroutine immediately, use any delay.
+	// To avoid running goroutine immediately and let the controller subscribe first,
+	// use any delay.
 	go func() {
 		time.Sleep(1 * time.Millisecond)
 
@@ -94,33 +122,6 @@ func TestController_watchAccountsChanges(t *testing.T) {
 			Accounts: []common.Address{address},
 		})
 	}()
-
-	c.startAccountWatcher([]uint64{chainID})
-
-	// Wait for event
-	<-ch
-
-	// Wait for DB to be cleaned up
-	c.accWatcher.Stop()
-
-	// Check that transfers, blocks and block ranges were deleted
-	transfers, err := database.GetTransfersByAddress(chainID, address, big.NewInt(2), 1)
-	require.NoError(t, err)
-	require.Len(t, transfers, 0)
-
-	blocksDAO := &BlockDAO{walletDB}
-	block, err := blocksDAO.GetLastBlockByAddress(chainID, address, 1)
-	require.NoError(t, err)
-	require.Nil(t, block)
-
-	ranges, _, err = blockRangesDAO.getBlockRange(chainID, address)
-	require.NoError(t, err)
-	require.Nil(t, ranges.eth.FirstKnown)
-	require.Nil(t, ranges.eth.LastKnown)
-	require.Nil(t, ranges.eth.Start)
-	require.Nil(t, ranges.tokens.FirstKnown)
-	require.Nil(t, ranges.tokens.LastKnown)
-	require.Nil(t, ranges.tokens.Start)
 }
 
 func TestController_cleanupAccountLeftovers(t *testing.T) {