diff options
| -rw-r--r-- | p2p/discover/table.go | 8 | ||||
| -rw-r--r-- | p2p/discover/udp_test.go | 55 |
2 files changed, 30 insertions, 33 deletions
diff --git a/p2p/discover/table.go b/p2p/discover/table.go index 4b7ddb775..f71320425 100644 --- a/p2p/discover/table.go +++ b/p2p/discover/table.go @@ -40,6 +40,8 @@ type Table struct { bonding map[NodeID]*bondproc bondslots chan struct{} // limits total number of active bonding processes + nodeAddedHook func(*Node) // for testing + net transport self *Node // metadata of the local node } @@ -431,6 +433,9 @@ func (tab *Table) pingreplace(new *Node, b *bucket) { } copy(b.entries[1:], b.entries) b.entries[0] = new + if tab.nodeAddedHook != nil { + tab.nodeAddedHook(new) + } } // ping a remote endpoint and wait for a reply, also updating the node database @@ -466,6 +471,9 @@ outer: } if len(bucket.entries) < bucketSize { bucket.entries = append(bucket.entries, n) + if tab.nodeAddedHook != nil { + tab.nodeAddedHook(n) + } } } } diff --git a/p2p/discover/udp_test.go b/p2p/discover/udp_test.go index 11fa31d7c..b5d035a98 100644 --- a/p2p/discover/udp_test.go +++ b/p2p/discover/udp_test.go @@ -234,14 +234,12 @@ func TestUDP_findnodeMultiReply(t *testing.T) { func TestUDP_successfulPing(t *testing.T) { test := newUDPTest(t) + added := make(chan *Node, 1) + test.table.nodeAddedHook = func(n *Node) { added <- n } defer test.table.Close() - done := make(chan struct{}) - go func() { - // The remote side sends a ping packet to initiate the exchange. - test.packetIn(nil, pingPacket, &ping{From: testRemote, To: testLocalAnnounced, Version: Version, Expiration: futureExp}) - close(done) - }() + // The remote side sends a ping packet to initiate the exchange. + go test.packetIn(nil, pingPacket, &ping{From: testRemote, To: testLocalAnnounced, Version: Version, Expiration: futureExp}) // the ping is replied to. test.waitPacketOut(func(p *pong) { @@ -277,35 +275,26 @@ func TestUDP_successfulPing(t *testing.T) { }) test.packetIn(nil, pongPacket, &pong{Expiration: futureExp}) - // ping should return shortly after getting the pong packet. - <-done - - // check that the node was added. - rid := PubkeyID(&test.remotekey.PublicKey) - rnode := find(test.table, rid) - if rnode == nil { - t.Fatalf("node %v not found in table", rid) - } - if !bytes.Equal(rnode.IP, test.remoteaddr.IP) { - t.Errorf("node has wrong IP: got %v, want: %v", rnode.IP, test.remoteaddr.IP) - } - if int(rnode.UDP) != test.remoteaddr.Port { - t.Errorf("node has wrong UDP port: got %v, want: %v", rnode.UDP, test.remoteaddr.Port) - } - if rnode.TCP != testRemote.TCP { - t.Errorf("node has wrong TCP port: got %v, want: %v", rnode.TCP, testRemote.TCP) - } -} - -func find(tab *Table, id NodeID) *Node { - for _, b := range tab.buckets { - for _, e := range b.entries { - if e.ID == id { - return e - } + // the node should be added to the table shortly after getting the + // pong packet. + select { + case n := <-added: + rid := PubkeyID(&test.remotekey.PublicKey) + if n.ID != rid { + t.Errorf("node has wrong ID: got %v, want %v", n.ID, rid) } + if !bytes.Equal(n.IP, test.remoteaddr.IP) { + t.Errorf("node has wrong IP: got %v, want: %v", n.IP, test.remoteaddr.IP) + } + if int(n.UDP) != test.remoteaddr.Port { + t.Errorf("node has wrong UDP port: got %v, want: %v", n.UDP, test.remoteaddr.Port) + } + if n.TCP != testRemote.TCP { + t.Errorf("node has wrong TCP port: got %v, want: %v", n.TCP, testRemote.TCP) + } + case <-time.After(2 * time.Second): + t.Errorf("node was not added within 2 seconds") } - return nil } // dgramPipe is a fake UDP socket. It queues all sent datagrams. |