From 8c701bf3b9e0bfd2e59ade988ef78982d12c5865 Mon Sep 17 00:00:00 2001
From: james-prysm <james@prysmaticlabs.com>
Date: Tue, 4 Nov 2025 22:52:49 -0600
Subject: [PATCH] fixing discovery tests

---
 beacon-chain/p2p/discovery_test.go | 64 ++++++++++++++++++++++++++++--
 1 file changed, 60 insertions(+), 4 deletions(-)

diff --git a/beacon-chain/p2p/discovery_test.go b/beacon-chain/p2p/discovery_test.go
index cd48877690..467467e54c 100644
--- a/beacon-chain/p2p/discovery_test.go
+++ b/beacon-chain/p2p/discovery_test.go
@@ -138,11 +138,18 @@ func setNodeSubnets(localNode *enode.LocalNode, attSubnets []uint64) {
 func TestCreateListener(t *testing.T) {
 	port := 1024
 	ipAddr, pkey := createAddrAndPrivKey(t)
+
+	db := testDB.SetupDB(t)
+	custodyInfoSet := make(chan struct{})
+	close(custodyInfoSet)
+
 	s := &Service{
+		ctx:                   t.Context(),
 		genesisTime:           time.Now(),
 		genesisValidatorsRoot: bytesutil.PadTo([]byte{'A'}, 32),
-		cfg:                   &Config{UDPPort: uint(port)},
+		cfg:                   &Config{UDPPort: uint(port), DB: db},
 		custodyInfo:           &custodyInfo{},
+		custodyInfoSet:        custodyInfoSet,
 	}
 	listener, err := s.createListener(ipAddr, pkey)
 	require.NoError(t, err)
@@ -165,11 +172,18 @@ func TestStartDiscV5_DiscoverAllPeers(t *testing.T) {
 	ipAddr, pkey := createAddrAndPrivKey(t)
 	genesisTime := time.Now()
 	genesisValidatorsRoot := make([]byte, 32)
+
+	db := testDB.SetupDB(t)
+	custodyInfoSet := make(chan struct{})
+	close(custodyInfoSet)
+
 	s := &Service{
-		cfg:                   &Config{UDPPort: uint(port), PingInterval: testPingInterval, DisableLivenessCheck: true},
+		ctx:                   t.Context(),
+		cfg:                   &Config{UDPPort: uint(port), PingInterval: testPingInterval, DisableLivenessCheck: true, DB: db},
 		genesisTime:           genesisTime,
 		genesisValidatorsRoot: genesisValidatorsRoot,
 		custodyInfo:           &custodyInfo{},
+		custodyInfoSet:        custodyInfoSet,
 	}
 	bootListener, err := s.createListener(ipAddr, pkey)
 	require.NoError(t, err)
@@ -189,13 +203,20 @@ func TestStartDiscV5_DiscoverAllPeers(t *testing.T) {
 			UDPPort:              uint(port),
 			PingInterval:         testPingInterval,
 			DisableLivenessCheck: true,
+			DB:                   db,
 		}
 		ipAddr, pkey := createAddrAndPrivKey(t)
+
+		custodyInfoSetLoop := make(chan struct{})
+		close(custodyInfoSetLoop)
+
 		s = &Service{
+			ctx:                   t.Context(),
 			cfg:                   cfg,
 			genesisTime:           genesisTime,
 			genesisValidatorsRoot: genesisValidatorsRoot,
 			custodyInfo:           &custodyInfo{},
+			custodyInfoSet:        custodyInfoSetLoop,
 		}
 		listener, err := s.startDiscoveryV5(ipAddr, pkey)
 		assert.NoError(t, err, "Could not start discovery for node")
@@ -346,11 +367,18 @@ func TestCreateLocalNode(t *testing.T) {
 func TestRebootDiscoveryListener(t *testing.T) {
 	port := 1024
 	ipAddr, pkey := createAddrAndPrivKey(t)
+
+	db := testDB.SetupDB(t)
+	custodyInfoSet := make(chan struct{})
+	close(custodyInfoSet)
+
 	s := &Service{
+		ctx:                   t.Context(),
 		genesisTime:           time.Now(),
 		genesisValidatorsRoot: bytesutil.PadTo([]byte{'A'}, 32),
-		cfg:                   &Config{UDPPort: uint(port)},
+		cfg:                   &Config{UDPPort: uint(port), DB: db},
 		custodyInfo:           &custodyInfo{},
+		custodyInfoSet:        custodyInfoSet,
 	}
 
 	createListener := func() (*discover.UDPv5, error) {
@@ -379,11 +407,17 @@ func TestRebootDiscoveryListener(t *testing.T) {
 func TestMultiAddrsConversion_InvalidIPAddr(t *testing.T) {
 	addr := net.ParseIP("invalidIP")
 	_, pkey := createAddrAndPrivKey(t)
+
+	custodyInfoSet := make(chan struct{})
+	close(custodyInfoSet)
+
 	s := &Service{
+		ctx:                   t.Context(),
 		genesisTime:           time.Now(),
 		genesisValidatorsRoot: bytesutil.PadTo([]byte{'A'}, 32),
 		cfg:                   &Config{},
 		custodyInfo:           &custodyInfo{},
+		custodyInfoSet:        custodyInfoSet,
 	}
 	node, err := s.createLocalNode(pkey, addr, 0, 0, 0)
 	require.NoError(t, err)
@@ -394,15 +428,23 @@ func TestMultiAddrsConversion_InvalidIPAddr(t *testing.T) {
 func TestMultiAddrConversion_OK(t *testing.T) {
 	hook := logTest.NewGlobal()
 	ipAddr, pkey := createAddrAndPrivKey(t)
+
+	db := testDB.SetupDB(t)
+	custodyInfoSet := make(chan struct{})
+	close(custodyInfoSet)
+
 	s := &Service{
+		ctx: t.Context(),
 		cfg: &Config{
 			UDPPort:  2000,
 			TCPPort:  3000,
 			QUICPort: 3000,
+			DB:       db,
 		},
 		genesisTime:           time.Now(),
 		genesisValidatorsRoot: bytesutil.PadTo([]byte{'A'}, 32),
 		custodyInfo:           &custodyInfo{},
+		custodyInfoSet:        custodyInfoSet,
 	}
 	listener, err := s.createListener(ipAddr, pkey)
 	require.NoError(t, err)
@@ -472,13 +514,20 @@ func TestHostIsResolved(t *testing.T) {
 		"2001:4860:4860::8844": true,
 	}
 
+	db := testDB.SetupDB(t)
+	custodyInfoSet := make(chan struct{})
+	close(custodyInfoSet)
+
 	s := &Service{
+		ctx: t.Context(),
 		cfg: &Config{
 			HostDNS: host,
+			DB:      db,
 		},
 		genesisTime:           time.Now(),
 		genesisValidatorsRoot: bytesutil.PadTo([]byte{'A'}, 32),
 		custodyInfo:           &custodyInfo{},
+		custodyInfoSet:        custodyInfoSet,
 	}
 	ip, key := createAddrAndPrivKey(t)
 	list, err := s.createListener(ip, key)
@@ -544,11 +593,18 @@ func TestUDPMultiAddress(t *testing.T) {
 	ipAddr, pkey := createAddrAndPrivKey(t)
 	genesisTime := time.Now()
 	genesisValidatorsRoot := make([]byte, 32)
+
+	db := testDB.SetupDB(t)
+	custodyInfoSet := make(chan struct{})
+	close(custodyInfoSet)
+
 	s := &Service{
-		cfg:                   &Config{UDPPort: uint(port)},
+		ctx:                   t.Context(),
+		cfg:                   &Config{UDPPort: uint(port), DB: db},
 		genesisTime:           genesisTime,
 		genesisValidatorsRoot: genesisValidatorsRoot,
 		custodyInfo:           &custodyInfo{},
+		custodyInfoSet:        custodyInfoSet,
 	}
 
 	createListener := func() (*discover.UDPv5, error) {