fix

Author: hero5512

p2p/discover/v4_udp_test.go

@@ -110,29 +110,50 @@ func (test *udpTest) packetInFrom(wantError error, key *ecdsa.PrivateKey, addr n
 
 // waits for a packet to be sent by the transport.
 // validate should have type func(X, netip.AddrPort, []byte), where X is a packet type.
-func (test *udpTest) waitPacketOut(validate interface{}) (closed bool) {
+// skipTypes is an optional list of packet types to skip
+func (test *udpTest) waitPacketOut(validate interface{}, skipTypes ...reflect.Type) (closed bool) {
 	test.t.Helper()
 
-	dgram, err := test.pipe.receive()
-	if err == errClosed {
-		return true
-	} else if err != nil {
-		test.t.Error("packet receive error:", err)
-		return false
-	}
-	p, _, hash, err := v4wire.Decode(dgram.data)
-	if err != nil {
-		test.t.Errorf("sent packet decode error: %v", err)
-		return false
-	}
-	fn := reflect.ValueOf(validate)
-	exptype := fn.Type().In(0)
-	if !reflect.TypeOf(p).AssignableTo(exptype) {
-		test.t.Errorf("sent packet type mismatch, got: %v, want: %v", reflect.TypeOf(p), exptype)
+	timeout := time.After(1 * time.Minute)
+	for {
+		select {
+		case <-timeout:
+			test.t.Fatal("timeout waiting for packet")
+		default:
+		}
+
+		dgram, err := test.pipe.receive()
+		if err == errClosed {
+			return true
+		} else if err != nil {
+			test.t.Error("packet receive error:", err)
+			return false
+		}
+		p, _, hash, err := v4wire.Decode(dgram.data)
+		if err != nil {
+			test.t.Errorf("sent packet decode error: %v", err)
+			return false
+		}
+		ptype := reflect.TypeOf(p)
+		skip := false
+		for _, skipType := range skipTypes {
+			if ptype == skipType {
+				skip = true
+				break
+			}
+		}
+		if skip {
+			continue
+		}
+		fn := reflect.ValueOf(validate)
+		exptype := fn.Type().In(0)
+		if !reflect.TypeOf(p).AssignableTo(exptype) {
+			test.t.Errorf("sent packet type mismatch, got: %v, want: %v", reflect.TypeOf(p), exptype)
+			return false
+		}
+		fn.Call([]reflect.Value{reflect.ValueOf(p), reflect.ValueOf(dgram.to), reflect.ValueOf(hash)})
 		return false
 	}
-	fn.Call([]reflect.Value{reflect.ValueOf(p), reflect.ValueOf(dgram.to), reflect.ValueOf(hash)})
-	return false
 }
 
 func TestUDPv4_packetErrors(t *testing.T) {
@@ -282,46 +303,20 @@ func TestUDPv4_findnode(t *testing.T) {
 	expected := test.table.findnodeByID(testTarget.ID(), bucketSize, true)
 	test.packetIn(nil, &v4wire.Findnode{Target: testTarget, Expiration: futureExp})
 	waitNeighbors := func(want []*enode.Node) {
-		timeout := time.After(1 * time.Minute)
-		for {
-			select {
-			case <-timeout:
-				t.Fatal("timeout waiting for neighbors response")
-			default:
-			}
-			dgram, err := test.pipe.receive()
-			if err == errClosed {
-				t.Fatal("socket closed before receiving neighbors")
-			} else if err != nil {
-				t.Fatal("packet receive error:", err)
-			}
-			p, _, _, err := v4wire.Decode(dgram.data)
-			if err != nil {
-				t.Fatalf("sent packet decode error: %v", err)
-			}
-			// Skip any PING packets from revalidation
-			if _, ok := p.(*v4wire.Ping); ok {
-				continue
-			}
-			// Check we got NEIGHBORS
-			neighbors, ok := p.(*v4wire.Neighbors)
-			if !ok {
-				t.Fatalf("sent packet type mismatch, got: %T, want: *v4wire.Neighbors", p)
-			}
-			if len(neighbors.Nodes) != len(want) {
-				t.Errorf("wrong number of results: got %d, want %d", len(neighbors.Nodes), len(want))
+		test.waitPacketOut(func(p *v4wire.Neighbors, to netip.AddrPort, hash []byte) {
+			if len(p.Nodes) != len(want) {
+				t.Errorf("wrong number of results: got %d, want %d", len(p.Nodes), len(want))
 				return
 			}
-			for i, n := range neighbors.Nodes {
+			for i, n := range p.Nodes {
 				if n.ID.ID() != want[i].ID() {
 					t.Errorf("result mismatch at %d:\n  got: %v\n  want: %v", i, n, expected.entries[i])
 				}
 				if !live[n.ID.ID()] {
 					t.Errorf("result includes dead node %v", n.ID.ID())
 				}
 			}
-			return
-		}
+		}, reflect.TypeOf((*v4wire.Ping)(nil)))
 	}
 	// Receive replies.
 	want := expected.entries