fix flaky tests

Author: marten-seemann

autonat_test.go

@@ -21,34 +21,44 @@ import (
 // these are mock service implementations for testing
 func makeAutoNATServicePrivate(ctx context.Context, t *testing.T) host.Host {
 	h := bhost.NewBlankHost(swarmt.GenSwarm(t))
-	h.SetStreamHandler(AutoNATProto, sayAutoNATPrivate)
+	h.SetStreamHandler(AutoNATProto, sayPrivateStreamHandler(t))
 	return h
 }
 
-func makeAutoNATServicePublic(ctx context.Context, t *testing.T) host.Host {
-	h := bhost.NewBlankHost(swarmt.GenSwarm(t))
-	h.SetStreamHandler(AutoNATProto, sayAutoNATPublic)
-	return h
-}
-
-func sayAutoNATPrivate(s network.Stream) {
-	defer s.Close()
-	w := protoio.NewDelimitedWriter(s)
-	res := pb.Message{
-		Type:         pb.Message_DIAL_RESPONSE.Enum(),
-		DialResponse: newDialResponseError(pb.Message_E_DIAL_ERROR, "no dialable addresses"),
+func sayPrivateStreamHandler(t *testing.T) network.StreamHandler {
+	return func(s network.Stream) {
+		defer s.Close()
+		r := protoio.NewDelimitedReader(s, network.MessageSizeMax)
+		if err := r.ReadMsg(&pb.Message{}); err != nil {
+			t.Error(err)
+			return
+		}
+		w := protoio.NewDelimitedWriter(s)
+		res := pb.Message{
+			Type:         pb.Message_DIAL_RESPONSE.Enum(),
+			DialResponse: newDialResponseError(pb.Message_E_DIAL_ERROR, "no dialable addresses"),
+		}
+		w.WriteMsg(&res)
 	}
-	w.WriteMsg(&res)
 }
 
-func sayAutoNATPublic(s network.Stream) {
-	defer s.Close()
-	w := protoio.NewDelimitedWriter(s)
-	res := pb.Message{
-		Type:         pb.Message_DIAL_RESPONSE.Enum(),
-		DialResponse: newDialResponseOK(s.Conn().RemoteMultiaddr()),
-	}
-	w.WriteMsg(&res)
+func makeAutoNATServicePublic(ctx context.Context, t *testing.T) host.Host {
+	h := bhost.NewBlankHost(swarmt.GenSwarm(t))
+	h.SetStreamHandler(AutoNATProto, func(s network.Stream) {
+		defer s.Close()
+		r := protoio.NewDelimitedReader(s, network.MessageSizeMax)
+		if err := r.ReadMsg(&pb.Message{}); err != nil {
+			t.Error(err)
+			return
+		}
+		w := protoio.NewDelimitedWriter(s)
+		res := pb.Message{
+			Type:         pb.Message_DIAL_RESPONSE.Enum(),
+			DialResponse: newDialResponseOK(s.Conn().RemoteMultiaddr()),
+		}
+		w.WriteMsg(&res)
+	})
+	return h
 }
 
 func makeAutoNAT(ctx context.Context, t *testing.T, ash host.Host) (host.Host, AutoNAT) {
@@ -184,7 +194,7 @@ func TestAutoNATPublictoPrivate(t *testing.T) {
 
 	expectEvent(t, s, network.ReachabilityPublic)
 
-	hs.SetStreamHandler(AutoNATProto, sayAutoNATPrivate)
+	hs.SetStreamHandler(AutoNATProto, sayPrivateStreamHandler(t))
 	hps := makeAutoNATServicePrivate(ctx, t)
 	connect(t, hps, hc)
 	identifyAsServer(hps, hc)
@@ -252,7 +262,7 @@ func TestAutoNATObservationRecording(t *testing.T) {
 	case <-s.Out():
 		t.Fatal("not expecting a public reachability event")
 	default:
-		//expected
+		// expected
 	}
 
 	addr, _ := ma.NewMultiaddr("/ip4/127.0.0.1/udp/1234")