diff --git a/dkg/dkg_test.go b/dkg/dkg_test.go index 657fbcb58..14e49df54 100644 --- a/dkg/dkg_test.go +++ b/dkg/dkg_test.go @@ -21,6 +21,7 @@ import ( eth2p0 "github.com/attestantio/go-eth2-client/spec/phase0" k1 "github.com/decred/dcrd/dcrec/secp256k1/v4" + "github.com/libp2p/go-libp2p/core/host" "github.com/libp2p/go-libp2p/core/peer" "github.com/stretchr/testify/require" "golang.org/x/sync/errgroup" @@ -580,11 +581,17 @@ func TestSyncFlow(t *testing.T) { stopDkgs = make([]context.CancelFunc, test.nodes) cTracker = newConnTracker(pIDs) dkgErrChan = make(chan error, test.nodes) // Buffered to prevent blocking + hosts = make([]host.Host, test.nodes) ) // Set all callbacks before starting any DKG processes to avoid race conditions for i := range test.nodes { configs[i].TestConfig.SyncCallback = cTracker.Set + + idx := i + configs[i].TestConfig.P2PNodeCallback = func(h host.Host) { + hosts[idx] = h + } } // Start DKG for initial peers. @@ -606,6 +613,15 @@ func TestSyncFlow(t *testing.T) { log.Info(ctx, "Stopping peer", z.Int("peer_index", idx)) stopDkgs[idx]() + // Manually close connections from other peers to this peer + for i, h := range hosts { + // The hosts array may contain nil values for peers that have been stopped. + // This safety check ensures we only attempt to close connections for active hosts. + if h != nil && i != idx { + _ = h.Network().ClosePeer(pIDs[idx]) + } + } + // Wait for this dkg process to return. select { case err := <-dkgErrChan: @@ -615,9 +631,6 @@ func TestSyncFlow(t *testing.T) { } } - // Give remaining peers time to detect disconnections and update their connection counts - time.Sleep(500 * time.Millisecond) - // Wait for remaining-initial peers to update connection counts. expect = len(test.connect) - len(test.disconnect) - 1 for _, idx := range test.connect {