From 952f7cbc5cf5d3e16d9171f1151b20c2a227d484 Mon Sep 17 00:00:00 2001 From: Sukun Date: Tue, 11 Apr 2023 10:35:47 +0530 Subject: [PATCH] autonat: fix flaky TestAutoNATDialRefused (#2245) --- p2p/host/autonat/autonat_test.go | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/p2p/host/autonat/autonat_test.go b/p2p/host/autonat/autonat_test.go index 617f6012a..d8cfcc51e 100644 --- a/p2p/host/autonat/autonat_test.go +++ b/p2p/host/autonat/autonat_test.go @@ -42,17 +42,25 @@ func sayPrivateStreamHandler(t *testing.T) network.StreamHandler { } } -func makeAutoNATRefuseDialRequest(t *testing.T) host.Host { +func makeAutoNATRefuseDialRequest(t *testing.T, done chan struct{}) host.Host { h := bhost.NewBlankHost(swarmt.GenSwarm(t)) - h.SetStreamHandler(AutoNATProto, sayRefusedStreamHandler(t)) + h.SetStreamHandler(AutoNATProto, sayRefusedStreamHandler(t, done)) return h } -func sayRefusedStreamHandler(t *testing.T) network.StreamHandler { +func sayRefusedStreamHandler(t *testing.T, done chan struct{}) network.StreamHandler { return func(s network.Stream) { defer s.Close() r := pbio.NewDelimitedReader(s, network.MessageSizeMax) if err := r.ReadMsg(&pb.Message{}); err != nil { + // ignore error if the test has completed + select { + case _, ok := <-done: + if !ok { + return + } + default: + } t.Error(err) return } @@ -238,14 +246,16 @@ func TestAutoNATDialRefused(t *testing.T) { connect(t, hs, hc) expectEvent(t, s, network.ReachabilityPublic, 10*time.Second) - hs.SetStreamHandler(AutoNATProto, sayRefusedStreamHandler(t)) - hps := makeAutoNATRefuseDialRequest(t) + done := make(chan struct{}) + hs.SetStreamHandler(AutoNATProto, sayRefusedStreamHandler(t, done)) + hps := makeAutoNATRefuseDialRequest(t, done) connect(t, hps, hc) identifyAsServer(hps, hc) require.Never(t, func() bool { return an.Status() != network.ReachabilityPublic }, 3*time.Second, 1*time.Second, "Expected probe to not change reachability from public") + close(done) } func TestAutoNATObservationRecording(t *testing.T) {