Fixed some races. Added more tests.

network/node_test.go

@@ -2,6 +2,8 @@ package network
 
 import (
 	"testing"
+
+	pb "github.com/micro/go-micro/network/proto"
 )
 
 var (
@@ -229,3 +231,111 @@ func TestTopology(t *testing.T) {
 		}
 	}
 }
+
+func TestUpdatePeerTopology(t *testing.T) {
+	// single node
+	single := &node{
+		id:      testNodeId,
+		address: testNodeAddress,
+		peers:   make(map[string]*node),
+		network: newNetwork(Name(testNodeNetName)),
+	}
+	// nil peer should return error
+	if err := single.updatePeerTopology(nil, 5); err == nil {
+		t.Errorf("Expected error, got %s", err)
+	}
+
+	// update with peer that is not yet in the peer map
+	pbPeer := &pb.Peer{
+		Node: &pb.Node{
+			Id:      "newPeer",
+			Address: "newPeerAddress",
+		},
+		Peers: make([]*pb.Peer, 0),
+	}
+	// it should add pbPeer to the single node peers
+	if err := single.updatePeerTopology(pbPeer, 5); err != nil {
+		t.Errorf("Error updating topology: %s", err)
+	}
+	if _, ok := single.peers[pbPeer.Node.Id]; !ok {
+		t.Errorf("Expected %s to be added to %s peers", pbPeer.Node.Id, single.id)
+	}
+
+	// complicated node graph
+	node := testSetup()
+	// build a simple topology to update node peer1
+	peer1 := node.peers["peer1"]
+	pbPeer1Node := &pb.Node{
+		Id:      peer1.id,
+		Address: peer1.address,
+	}
+
+	pbPeer111 := &pb.Peer{
+		Node: &pb.Node{
+			Id:      "peer111",
+			Address: "peer111Address",
+		},
+		Peers: make([]*pb.Peer, 0),
+	}
+
+	pbPeer121 := &pb.Peer{
+		Node: &pb.Node{
+			Id:      "peer121",
+			Address: "peer121Address",
+		},
+		Peers: make([]*pb.Peer, 0),
+	}
+	// topology to update
+	pbPeer1 := &pb.Peer{
+		Node:  pbPeer1Node,
+		Peers: []*pb.Peer{pbPeer111, pbPeer121},
+	}
+	// update peer1 topology
+	if err := node.updatePeerTopology(pbPeer1, 5); err != nil {
+		t.Errorf("Error updating topology: %s", err)
+	}
+	// make sure peer1 topology has been correctly updated
+	newPeerIds := []string{pbPeer111.Node.Id, pbPeer121.Node.Id}
+	for _, id := range newPeerIds {
+		if _, ok := node.peers["peer1"].peers[id]; !ok {
+			t.Errorf("Expected %s to be a peer of %s", id, "peer1")
+		}
+	}
+}
+
+func TestGetProtoTopology(t *testing.T) {
+	// single node
+	single := &node{
+		id:      testNodeId,
+		address: testNodeAddress,
+		peers:   make(map[string]*node),
+		network: newNetwork(Name(testNodeNetName)),
+	}
+	topCount := 0
+
+	protoTop, err := single.getProtoTopology(10)
+	if err != nil {
+		t.Errorf("Error getting proto topology: %s", err)
+	}
+	if len(protoTop.Peers) != topCount {
+		t.Errorf("Expected to find %d nodes, found: %d", topCount, len(protoTop.Peers))
+	}
+
+	// complicated node graph
+	node := testSetup()
+	topCount = 3
+	// list of ids of nodes of depth 1 i.e. node peers
+	peerIds := make(map[string]bool)
+	// add peer Ids
+	for _, id := range testNodePeerIds {
+		peerIds[id] = true
+	}
+	// depth 1 should give us immmediate neighbours only
+	protoTop, err = node.getProtoTopology(1)
+	if err != nil {
+		t.Errorf("Error getting proto topology: %s", err)
+	}
+	if len(protoTop.Peers) != topCount {
+		t.Errorf("Expected to find %d nodes, found: %d", topCount, len(protoTop.Peers))
+	}
+}