From 5a704d2daa5aa4a3afd70ea61f08d5c4482e3fe3 Mon Sep 17 00:00:00 2001
From: Cole Brown <bigswim@gmail.com>
Date: Thu, 28 Jun 2018 17:57:19 -0400
Subject: [PATCH] Add Stat to mock net conns and streams

---
 p2p/net/mock/mock_conn.go   | 9 ++++++++-
 p2p/net/mock/mock_link.go   | 8 ++++----
 p2p/net/mock/mock_stream.go | 8 +++++++-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/p2p/net/mock/mock_conn.go b/p2p/net/mock/mock_conn.go
index 68e9c84..f4cf103 100644
--- a/p2p/net/mock/mock_conn.go
+++ b/p2p/net/mock/mock_conn.go
@@ -29,14 +29,16 @@ type conn struct {
 	rconn   *conn // counterpart
 	streams list.List
 	proc    process.Process
+	stat    inet.Stat
 
 	sync.RWMutex
 }
 
-func newConn(ln, rn *peernet, l *link) *conn {
+func newConn(ln, rn *peernet, l *link, dir inet.Direction) *conn {
 	c := &conn{net: ln, link: l}
 	c.local = ln.peer
 	c.remote = rn.peer
+	c.stat = inet.Stat{Direction: dir}
 
 	c.localAddr = ln.ps.Addrs(ln.peer)[0]
 	c.remoteAddr = rn.ps.Addrs(rn.peer)[0]
@@ -155,3 +157,8 @@ func (c *conn) RemotePeer() peer.ID {
 func (c *conn) RemotePublicKey() ic.PubKey {
 	return c.remotePubKey
 }
+
+// Stat returns metadata about the connection
+func (c *conn) Stat() inet.Stat {
+	return c.stat
+}
diff --git a/p2p/net/mock/mock_link.go b/p2p/net/mock/mock_link.go
index d4d6053..9c6049a 100644
--- a/p2p/net/mock/mock_link.go
+++ b/p2p/net/mock/mock_link.go
@@ -33,8 +33,8 @@ func (l *link) newConnPair(dialer *peernet) (*conn, *conn) {
 	l.RLock()
 	defer l.RUnlock()
 
-	c1 := newConn(l.nets[0], l.nets[1], l)
-	c2 := newConn(l.nets[1], l.nets[0], l)
+	c1 := newConn(l.nets[0], l.nets[1], l, inet.DirOutbound)
+	c2 := newConn(l.nets[1], l.nets[0], l, inet.DirInbound)
 	c1.rconn = c2
 	c2.rconn = c1
 
@@ -48,8 +48,8 @@ func (l *link) newStreamPair() (*stream, *stream) {
 	ra, wb := io.Pipe()
 	rb, wa := io.Pipe()
 
-	sa := NewStream(wa, ra)
-	sb := NewStream(wb, rb)
+	sa := NewStream(wa, ra, inet.DirOutbound)
+	sb := NewStream(wb, rb, inet.DirInbound)
 	return sa, sb
 }
 
diff --git a/p2p/net/mock/mock_stream.go b/p2p/net/mock/mock_stream.go
index 6d056c8..4f0c395 100644
--- a/p2p/net/mock/mock_stream.go
+++ b/p2p/net/mock/mock_stream.go
@@ -25,6 +25,7 @@ type stream struct {
 	writeErr error
 
 	protocol protocol.ID
+	stat     inet.Stat
 }
 
 var ErrReset error = errors.New("stream reset")
@@ -35,7 +36,7 @@ type transportObject struct {
 	arrivalTime time.Time
 }
 
-func NewStream(w *io.PipeWriter, r *io.PipeReader) *stream {
+func NewStream(w *io.PipeWriter, r *io.PipeReader, dir inet.Direction) *stream {
 	s := &stream{
 		read:      r,
 		write:     w,
@@ -43,6 +44,7 @@ func NewStream(w *io.PipeWriter, r *io.PipeReader) *stream {
 		close:     make(chan struct{}, 1),
 		closed:    make(chan struct{}),
 		toDeliver: make(chan *transportObject),
+		stat:      inet.Stat{Direction: dir},
 	}
 
 	go s.transport()
@@ -66,6 +68,10 @@ func (s *stream) Protocol() protocol.ID {
 	return s.protocol
 }
 
+func (s *stream) Stat() inet.Stat {
+	return s.stat
+}
+
 func (s *stream) SetProtocol(proto protocol.ID) {
 	s.protocol = proto
 }
-- 
GitLab