yggdrasil-go/src/yggdrasil/core_test.go

package yggdrasil

import (
	"bytes"
	"math/rand"
	"os"
	"testing"
	"time"

	"github.com/gologme/log"

	"github.com/yggdrasil-network/yggdrasil-go/src/config"
)

// GenerateConfig is modification
func GenerateConfig() *config.NodeConfig {
	cfg := config.GenerateConfig()
	cfg.AdminListen = "none"
	cfg.Listen = []string{"tcp://127.0.0.1:0"}
	cfg.IfName = "none"

	return cfg
}

func GetLoggerWithPrefix(prefix string, verbose bool) *log.Logger {
	l := log.New(os.Stderr, prefix, log.Flags())
	if !verbose {
		return l
	}
	l.EnableLevel("info")
	l.EnableLevel("warn")
	l.EnableLevel("error")
	return l
}

func CreateAndConnectTwo(t testing.TB, verbose bool) (*Core, *Core) {
	nodeA := Core{}
	_, err := nodeA.Start(GenerateConfig(), GetLoggerWithPrefix("A: ", verbose))
	if err != nil {
		t.Fatal(err)
	}

	nodeB := Core{}
	_, err = nodeB.Start(GenerateConfig(), GetLoggerWithPrefix("B: ", verbose))
	if err != nil {
		t.Fatal(err)
	}

	err = nodeB.AddPeer("tcp://"+nodeA.link.tcp.getAddr().String(), "")
	if err != nil {
		t.Fatal(err)
	}

	if l := len(nodeA.GetPeers()); l != 1 {
		t.Fatal("unexpected number of peers", l)
	}
	if l := len(nodeB.GetPeers()); l != 1 {
		t.Fatal("unexpected number of peers", l)
	}

	return &nodeA, &nodeB
}

// WaitConnected blocks until either nodes negotiated DHT or 5 seconds passed.
func WaitConnected(nodeA, nodeB *Core) bool {
	// It may take up to 3 seconds, but let's wait 5.
	for i := 0; i < 50; i++ {
		time.Sleep(100 * time.Millisecond)
		if len(nodeA.GetSwitchPeers()) > 0 && len(nodeB.GetSwitchPeers()) > 0 {
			return true
		}
	}
	return false
}

func TestCore_Start_Connect(t *testing.T) {
	CreateAndConnectTwo(t, true)
}

func CreateEchoListener(t testing.TB, nodeA *Core, bufLen int, repeats int) chan struct{} {
	// Listen
	listener, err := nodeA.ConnListen()
	if err != nil {
		t.Fatal(err)
	}

	done := make(chan struct{})
	go func() {
		defer listener.Close()
		conn, err := listener.Accept()
		if err != nil {
			t.Error(err)
			return
		}
		defer conn.Close()
		buf := make([]byte, bufLen)

		for i := 0; i < repeats; i++ {
			n, err := conn.Read(buf)
			if err != nil {
				t.Error(err)
				return
			}
			if n != bufLen {
				t.Error("missing data")
				return
			}
			_, err = conn.Write(buf)
			if err != nil {
				t.Error(err)
			}
		}
		done <- struct{}{}
	}()

	return done
}

func TestCore_Start_Transfer(t *testing.T) {
	nodeA, nodeB := CreateAndConnectTwo(t, true)

	msgLen := 1500
	done := CreateEchoListener(t, nodeA, msgLen, 1)

	if !WaitConnected(nodeA, nodeB) {
		t.Fatal("nodes did not connect")
	}

	// Dial
	dialer, err := nodeB.ConnDialer()
	if err != nil {
		t.Fatal(err)
	}
	conn, err := dialer.Dial("nodeid", nodeA.NodeID().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	msg := make([]byte, msgLen)
	rand.Read(msg)
	conn.Write(msg)
	if err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, msgLen)
	_, err = conn.Read(buf)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(msg, buf) != 0 {
		t.Fatal("expected echo")
	}
	<-done
}

func BenchmarkCore_Start_Transfer(b *testing.B) {
	nodeA, nodeB := CreateAndConnectTwo(b, false)

	msgLen := 1500 // typical MTU
	done := CreateEchoListener(b, nodeA, msgLen, b.N)

	if !WaitConnected(nodeA, nodeB) {
		b.Fatal("nodes did not connect")
	}

	// Dial
	dialer, err := nodeB.ConnDialer()
	if err != nil {
		b.Fatal(err)
	}
	conn, err := dialer.Dial("nodeid", nodeA.NodeID().String())
	if err != nil {
		b.Fatal(err)
	}
	defer conn.Close()
	msg := make([]byte, msgLen)
	rand.Read(msg)
	buf := make([]byte, msgLen)

	b.SetBytes(int64(b.N * msgLen))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		conn.Write(msg)
		if err != nil {
			b.Fatal(err)
		}
		_, err = conn.Read(buf)
		if err != nil {
			b.Fatal(err)
		}
	}
	<-done
}