ceremonyclient/go-libp2p/p2p/net/mock/mock_notif_test.go

package mocknet

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"

	ma "github.com/multiformats/go-multiaddr"
	"github.com/stretchr/testify/require"
)

func TestNotifications(t *testing.T) {
	const swarmSize = 5
	const timeout = 10 * time.Second

	mn, err := FullMeshLinked(swarmSize)
	if err != nil {
		t.Fatal(err)
	}
	defer mn.Close()

	// signup notifs
	nets := mn.Nets()
	notifiees := make(map[peer.ID]*netNotifiee, len(nets))
	for _, pn := range nets {
		defer pn.Close()

		n := newNetNotifiee(t, swarmSize)
		pn.Notify(n)
		notifiees[pn.LocalPeer()] = n
	}

	// connect all but self
	if err := mn.ConnectAllButSelf(); err != nil {
		t.Fatal(err)
	}

	// test everyone got the correct connection opened calls
	for _, s1 := range nets {
		n := notifiees[s1.LocalPeer()]
		notifs := make(map[peer.ID][]network.Conn)
		for _, s2 := range nets {
			if s2 == s1 {
				continue
			}

			// this feels a little sketchy, but its probably okay
			for len(s1.ConnsToPeer(s2.LocalPeer())) != len(notifs[s2.LocalPeer()]) {
				select {
				case c := <-n.connected:
					nfp := notifs[c.RemotePeer()]
					notifs[c.RemotePeer()] = append(nfp, c)
				case <-time.After(timeout):
					t.Fatal("timeout")
				}
			}
		}

		for p, cons := range notifs {
			expect := s1.ConnsToPeer(p)
			if len(expect) != len(cons) {
				t.Fatal("got different number of connections")
			}

			for _, c := range cons {
				var found bool
				for _, c2 := range expect {
					if c == c2 {
						found = true
						break
					}
				}

				if !found {
					t.Fatal("connection not found!")
				}
			}
		}
	}

	acceptedStream := make(chan struct{}, 1000)
	for _, s := range nets {
		s.SetStreamHandler(func(s network.Stream) {
			acceptedStream <- struct{}{}
			s.Close()
		})
	}

	// Make sure we've received at last one stream per conn.
	for _, s := range nets {
		conns := s.Conns()
		for _, c := range conns {
			st1, err := c.NewStream(context.Background())
			if err != nil {
				t.Error(err)
				continue
			}
			t.Logf("%s %s <--%p--> %s %s", c.LocalPeer(), c.LocalMultiaddr(), st1, c.RemotePeer(), c.RemoteMultiaddr())
			st1.Close()
		}
	}

	// close conns
	for _, s1 := range nets {
		n1 := notifiees[s1.LocalPeer()]
		for _, c1 := range s1.Conns() {
			c2 := ConnComplement(c1)

			n2 := notifiees[c2.LocalPeer()]
			c1.Close()

			var c3, c4 network.Conn
			select {
			case c3 = <-n1.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c1 != c3 {
				t.Fatal("got incorrect conn", c1, c3)
			}

			select {
			case c4 = <-n2.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c2 != c4 {
				t.Fatal("got incorrect conn", c1, c2)
			}
		}
	}

	for _, n1 := range notifiees {
		// Avoid holding this lock while waiting, otherwise we can deadlock.
		streamStateCopy := map[network.Stream]chan struct{}{}
		n1.streamState.Lock()
		for str, ch := range n1.streamState.m {
			streamStateCopy[str] = ch
		}
		n1.streamState.Unlock()

		for str1, ch1 := range streamStateCopy {
			<-ch1
			str2 := StreamComplement(str1)
			n2 := notifiees[str1.Conn().RemotePeer()]

			// make sure the OpenedStream notification was processed first
			var ch2 chan struct{}
			require.Eventually(t, func() bool {
				n2.streamState.Lock()
				defer n2.streamState.Unlock()
				ch, ok := n2.streamState.m[str2]
				if ok {
					ch2 = ch
				}
				return ok
			}, time.Second, 10*time.Millisecond)

			<-ch2
		}
	}
}

type netNotifiee struct {
	t *testing.T

	listen       chan ma.Multiaddr
	listenClose  chan ma.Multiaddr
	connected    chan network.Conn
	disconnected chan network.Conn

	streamState struct {
		sync.Mutex
		m map[network.Stream]chan struct{}
	}
}

func newNetNotifiee(t *testing.T, buffer int) *netNotifiee {
	nn := &netNotifiee{
		t:            t,
		listen:       make(chan ma.Multiaddr, 1),
		listenClose:  make(chan ma.Multiaddr, 1),
		connected:    make(chan network.Conn, buffer*2),
		disconnected: make(chan network.Conn, buffer*2),
	}
	nn.streamState.m = make(map[network.Stream]chan struct{})
	return nn
}

func (nn *netNotifiee) Listen(n network.Network, a ma.Multiaddr) {
	nn.listen <- a
}
func (nn *netNotifiee) ListenClose(n network.Network, a ma.Multiaddr) {
	nn.listenClose <- a
}
func (nn *netNotifiee) Connected(n network.Network, v network.Conn) {
	nn.connected <- v
}
func (nn *netNotifiee) Disconnected(n network.Network, v network.Conn) {
	nn.disconnected <- v
}