neo-go/pkg/peer/peer_test.go

package peer_test

import (
	"net"
	"testing"
	"time"

	"github.com/CityOfZion/neo-go/pkg/peer"
	"github.com/CityOfZion/neo-go/pkg/wire"
	"github.com/CityOfZion/neo-go/pkg/wire/payload"
	"github.com/CityOfZion/neo-go/pkg/wire/protocol"
	"github.com/stretchr/testify/assert"
)

func returnConfig() peer.LocalConfig {

	DefaultHeight := func() uint32 {
		return 10
	}

	OnAddr := func(p *peer.Peer, msg *payload.AddrMessage) {}
	OnHeader := func(p *peer.Peer, msg *payload.HeadersMessage) {}
	OnGetHeaders := func(p *peer.Peer, msg *payload.GetHeadersMessage) {}
	OnInv := func(p *peer.Peer, msg *payload.InvMessage) {}
	OnGetData := func(p *peer.Peer, msg *payload.GetDataMessage) {}
	OnBlock := func(p *peer.Peer, msg *payload.BlockMessage) {}
	OnGetBlocks := func(p *peer.Peer, msg *payload.GetBlocksMessage) {}

	return peer.LocalConfig{
		Net:         protocol.MainNet,
		UserAgent:   "NEO-GO-Default",
		Services:    protocol.NodePeerService,
		Nonce:       1200,
		ProtocolVer: 0,
		Relay:       false,
		Port:        10332,
		// pointer to config will keep the startheight updated for each version
		//Message we plan to send
		StartHeight:  DefaultHeight,
		OnHeader:     OnHeader,
		OnAddr:       OnAddr,
		OnGetHeaders: OnGetHeaders,
		OnInv:        OnInv,
		OnGetData:    OnGetData,
		OnBlock:      OnBlock,
		OnGetBlocks:  OnGetBlocks,
	}
}

func TestHandshake(t *testing.T) {
	address := ":20338"
	go func() {

		conn, err := net.DialTimeout("tcp", address, 2*time.Second)
		if err != nil {
			t.Fatal(err)
		}
		p := peer.NewPeer(conn, true, returnConfig())
		err = p.Run()
		verack, err := payload.NewVerackMessage()
		if err != nil {
			t.Fail()
		}
		if err := p.Write(verack); err != nil {
			t.Fatal(err)
		}

		assert.Equal(t, true, p.IsVerackReceived())

	}()

	listener, err := net.Listen("tcp", address)
	if err != nil {
		t.Fatal(err)
		return
	}

	defer func() {
		listener.Close()
	}()

	for {

		conn, err := listener.Accept()
		if err != nil {
			t.Fatal(err)
		}

		tcpAddrMe := &net.TCPAddr{IP: net.ParseIP("82.2.97.142"), Port: 20338}
		nonce := uint32(100)
		messageVer, err := payload.NewVersionMessage(tcpAddrMe, 2595770, false, protocol.DefaultVersion, protocol.UserAgent, nonce, protocol.NodePeerService)

		if err != nil {
			t.Fatal(err)
		}

		if err := wire.WriteMessage(conn, protocol.MainNet, messageVer); err != nil {
			t.Fatal(err)
			return
		}

		readmsg, err := wire.ReadMessage(conn, protocol.MainNet)
		if err != nil {
			t.Fatal(err)
		}
		version, ok := readmsg.(*payload.VersionMessage)
		if !ok {
			t.Fatal(err)
		}

		assert.NotEqual(t, nil, version)

		messageVrck, err := payload.NewVerackMessage()
		if err != nil {
			t.Fatal(err)
		}

		assert.NotEqual(t, nil, messageVrck)

		if err := wire.WriteMessage(conn, protocol.MainNet, messageVrck); err != nil {
			t.Fatal(err)
		}

		readmsg, err = wire.ReadMessage(conn, protocol.MainNet)
		if err != nil {
			t.Fatal(err)
		}

		assert.NotEqual(t, nil, readmsg)

		verk, ok := readmsg.(*payload.VerackMessage)
		if !ok {
			t.Fatal(err)
		}
		assert.NotEqual(t, nil, verk)

		return
	}

}

func TestConfigurations(t *testing.T) {
	_, conn := net.Pipe()

	inbound := true

	config := returnConfig()

	p := peer.NewPeer(conn, inbound, config)

	// test inbound
	assert.Equal(t, inbound, p.Inbound())

	// handshake not done, should be false
	assert.Equal(t, false, p.IsVerackReceived())

	assert.Equal(t, config.Services, p.Services())

	assert.Equal(t, config.Relay, p.CanRelay())

	assert.WithinDuration(t, time.Now(), p.CreatedAt(), 1*time.Second)
}

func TestPeerDisconnect(t *testing.T) {
	// Make sure everything is shutdown
	// Make sure timer is shutdown in stall detector too. Should maybe put this part of test into stall detector.

	_, conn := net.Pipe()
	inbound := true
	config := returnConfig()
	p := peer.NewPeer(conn, inbound, config)

	p.Disconnect()
	verack, err := payload.NewVerackMessage()
	assert.Nil(t, err)

	err = p.Write(verack)
	assert.NotNil(t, err)

	// Check if stall detector is still running
	_, ok := <-p.Detector.Quitch
	assert.Equal(t, ok, false)
}

func TestNotifyDisconnect(t *testing.T) {

	_, conn := net.Pipe()
	inbound := true
	config := returnConfig()
	p := peer.NewPeer(conn, inbound, config)

	p.Disconnect()
	p.NotifyDisconnect()
	// TestNotify uses default test timeout as the passing condition
	// Failure condition can be seen when you comment out p.Disconnect()
}