mocknet: use host

net2/mock/interface.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	ic "github.com/jbenet/go-ipfs/p2p/crypto"
+	host "github.com/jbenet/go-ipfs/p2p/host"
 	inet "github.com/jbenet/go-ipfs/p2p/net2"
 	peer "github.com/jbenet/go-ipfs/p2p/peer"
 
@@ -20,16 +21,18 @@ import (
 type Mocknet interface {
 
 	// GenPeer generates a peer and its inet.Network in the Mocknet
-	GenPeer() (inet.Network, error)
+	GenPeer() (host.Host, error)
 
 	// AddPeer adds an existing peer. we need both a privkey and addr.
 	// ID is derived from PrivKey
-	AddPeer(ic.PrivKey, ma.Multiaddr) (inet.Network, error)
+	AddPeer(ic.PrivKey, ma.Multiaddr) (host.Host, error)
 
 	// retrieve things (with randomized iteration order)
 	Peers() []peer.ID
 	Net(peer.ID) inet.Network
 	Nets() []inet.Network
+	Host(peer.ID) host.Host
+	Hosts() []host.Host
 	Links() LinkMap
 	LinksBetweenPeers(a, b peer.ID) []Link
 	LinksBetweenNets(a, b inet.Network) []Link

net2/mock/mock_net.go

@@ -5,6 +5,8 @@ import (
 	"sync"
 
 	ic "github.com/jbenet/go-ipfs/p2p/crypto"
+	host "github.com/jbenet/go-ipfs/p2p/host"
+	bhost "github.com/jbenet/go-ipfs/p2p/host/basic"
 	inet "github.com/jbenet/go-ipfs/p2p/net2"
 	peer "github.com/jbenet/go-ipfs/p2p/peer"
 	testutil "github.com/jbenet/go-ipfs/util/testutil"
@@ -16,9 +18,8 @@ import (
 
 // mocknet implements mocknet.Mocknet
 type mocknet struct {
-	// must map on peer.ID (instead of peer.ID) because
-	// each inet.Network has different peerstore
-	nets map[peer.ID]*peernet
+	nets  map[peer.ID]*peernet
+	hosts map[peer.ID]*bhost.BasicHost
 
 	// links make it possible to connect two peers.
 	// think of links as the physical medium.
@@ -35,12 +36,13 @@ type mocknet struct {
 func New(ctx context.Context) Mocknet {
 	return &mocknet{
 		nets:  map[peer.ID]*peernet{},
+		hosts: map[peer.ID]*bhost.BasicHost{},
 		links: map[peer.ID]map[peer.ID]map[*link]struct{}{},
 		cg:    ctxgroup.WithContext(ctx),
 	}
 }
 
-func (mn *mocknet) GenPeer() (inet.Network, error) {
+func (mn *mocknet) GenPeer() (host.Host, error) {
 	sk, _, err := testutil.RandKeyPair(512)
 	if err != nil {
 		return nil, err
@@ -48,20 +50,22 @@ func (mn *mocknet) GenPeer() (inet.Network, error) {
 
 	a := testutil.RandLocalTCPAddress()
 
-	n, err := mn.AddPeer(sk, a)
+	h, err := mn.AddPeer(sk, a)
 	if err != nil {
 		return nil, err
 	}
 
-	return n, nil
+	return h, nil
 }
 
-func (mn *mocknet) AddPeer(k ic.PrivKey, a ma.Multiaddr) (inet.Network, error) {
+func (mn *mocknet) AddPeer(k ic.PrivKey, a ma.Multiaddr) (host.Host, error) {
 	n, err := newPeernet(mn.cg.Context(), mn, k, a)
 	if err != nil {
 		return nil, err
 	}
 
+	h := bhost.New(n)
+
 	// make sure to add listening address!
 	// this makes debugging things simpler as remembering to register
 	// an address may cause unexpected failure.
@@ -72,8 +76,9 @@ func (mn *mocknet) AddPeer(k ic.PrivKey, a ma.Multiaddr) (inet.Network, error) {
 
 	mn.Lock()
 	mn.nets[n.peer] = n
+	mn.hosts[n.peer] = h
 	mn.Unlock()
-	return n, nil
+	return h, nil
 }
 
 func (mn *mocknet) Peers() []peer.ID {
@@ -87,16 +92,29 @@ func (mn *mocknet) Peers() []peer.ID {
 	return cp
 }
 
+func (mn *mocknet) Host(pid peer.ID) host.Host {
+	mn.RLock()
+	host := mn.hosts[pid]
+	mn.RUnlock()
+	return host
+}
+
 func (mn *mocknet) Net(pid peer.ID) inet.Network {
+	mn.RLock()
+	n := mn.nets[pid]
+	mn.RUnlock()
+	return n
+}
+
+func (mn *mocknet) Hosts() []host.Host {
 	mn.RLock()
 	defer mn.RUnlock()
 
-	for _, n := range mn.nets {
-		if n.peer == pid {
-			return n
-		}
+	cp := make([]host.Host, 0, len(mn.hosts))
+	for _, h := range mn.hosts {
+		cp = append(cp, h)
 	}
-	return nil
+	return cp
 }
 
 func (mn *mocknet) Nets() []inet.Network {

net2/mock/mock_test.go

@@ -9,6 +9,7 @@ import (
 
 	inet "github.com/jbenet/go-ipfs/p2p/net2"
 	peer "github.com/jbenet/go-ipfs/p2p/peer"
+	protocol "github.com/jbenet/go-ipfs/p2p/protocol"
 	testutil "github.com/jbenet/go-ipfs/util/testutil"
 
 	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
@@ -46,31 +47,44 @@ func TestNetworkSetup(t *testing.T) {
 	a2 := testutil.RandLocalTCPAddress()
 	a3 := testutil.RandLocalTCPAddress()
 
-	n1, err := mn.AddPeer(sk1, a1)
+	h1, err := mn.AddPeer(sk1, a1)
 	if err != nil {
 		t.Fatal(err)
 	}
-	p1 := n1.LocalPeer()
+	p1 := h1.ID()
 
-	n2, err := mn.AddPeer(sk2, a2)
+	h2, err := mn.AddPeer(sk2, a2)
 	if err != nil {
 		t.Fatal(err)
 	}
-	p2 := n2.LocalPeer()
+	p2 := h2.ID()
 
-	n3, err := mn.AddPeer(sk3, a3)
+	h3, err := mn.AddPeer(sk3, a3)
 	if err != nil {
 		t.Fatal(err)
 	}
-	p3 := n3.LocalPeer()
+	p3 := h3.ID()
 
 	// check peers and net
+	if mn.Host(p1) != h1 {
+		t.Error("host for p1.ID != h1")
+	}
+	if mn.Host(p2) != h2 {
+		t.Error("host for p2.ID != h2")
+	}
+	if mn.Host(p3) != h3 {
+		t.Error("host for p3.ID != h3")
+	}
+
+	n1 := h1.Network()
 	if mn.Net(p1) != n1 {
 		t.Error("net for p1.ID != n1")
 	}
+	n2 := h2.Network()
 	if mn.Net(p2) != n2 {
 		t.Error("net for p2.ID != n1")
 	}
+	n3 := h3.Network()
 	if mn.Net(p3) != n3 {
 		t.Error("net for p3.ID != n1")
 	}
@@ -275,12 +289,12 @@ func TestStreams(t *testing.T) {
 		s.Close()
 	}
 
-	nets := mn.Nets()
-	for _, n := range nets {
-		n.SetStreamHandler(handler)
+	hosts := mn.Hosts()
+	for _, h := range mn.Hosts() {
+		h.SetStreamHandler(protocol.TestingID, handler)
 	}
 
-	s, err := nets[0].NewStream(nets[1].LocalPeer())
+	s, err := hosts[0].NewStream(protocol.TestingID, hosts[1].ID())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -350,9 +364,10 @@ func TestStreamsStress(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	nets := mn.Nets()
-	for _, n := range nets {
-		n.SetStreamHandler(makePonger("pingpong"))
+	hosts := mn.Hosts()
+	for _, h := range hosts {
+		ponger := makePonger(string(protocol.TestingID))
+		h.SetStreamHandler(protocol.TestingID, ponger)
 	}
 
 	var wg sync.WaitGroup
@@ -360,11 +375,11 @@ func TestStreamsStress(t *testing.T) {
 		wg.Add(1)
 		go func(i int) {
 			defer wg.Done()
-			from := rand.Intn(len(nets))
-			to := rand.Intn(len(nets))
-			s, err := nets[from].NewStream(nets[to].LocalPeer())
+			from := rand.Intn(len(hosts))
+			to := rand.Intn(len(hosts))
+			s, err := hosts[from].NewStream(protocol.TestingID, hosts[to].ID())
 			if err != nil {
-				log.Debugf("%d (%s) %d (%s)", from, nets[from], to, nets[to])
+				log.Debugf("%d (%s) %d (%s)", from, hosts[from], to, hosts[to])
 				panic(err)
 			}
 
@@ -389,12 +404,12 @@ func TestAdding(t *testing.T) {
 		}
 
 		a := testutil.RandLocalTCPAddress()
-		n, err := mn.AddPeer(sk, a)
+		h, err := mn.AddPeer(sk, a)
 		if err != nil {
 			t.Fatal(err)
 		}
 
-		peers = append(peers, n.LocalPeer())
+		peers = append(peers, h.ID())
 	}
 
 	p1 := peers[0]
@@ -410,11 +425,11 @@ func TestAdding(t *testing.T) {
 	}
 
 	// set the new stream handler on p2
-	n2 := mn.Net(p2)
-	if n2 == nil {
-		t.Fatalf("no network for %s", p2)
+	h2 := mn.Host(p2)
+	if h2 == nil {
+		t.Fatalf("no host for %s", p2)
 	}
-	n2.SetStreamHandler(func(s inet.Stream) {
+	h2.SetStreamHandler(protocol.TestingID, func(s inet.Stream) {
 		defer s.Close()
 
 		b := make([]byte, 4)
@@ -436,12 +451,12 @@ func TestAdding(t *testing.T) {
 	}
 
 	// talk to p2
-	n1 := mn.Net(p1)
-	if n1 == nil {
+	h1 := mn.Host(p1)
+	if h1 == nil {
 		t.Fatalf("no network for %s", p1)
 	}
 
-	s, err := n1.NewStream(p2)
+	s, err := h1.NewStream(protocol.TestingID, p2)
 	if err != nil {
 		t.Fatal(err)
 	}