integration test preliminaries

Author: vyzo

itest/echo.go

@@ -0,0 +1,201 @@
+package itest
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"sync"
+	"time"
+
+	"github.com/libp2p/go-libp2p-core/host"
+	"github.com/libp2p/go-libp2p-core/network"
+	"github.com/libp2p/go-libp2p-core/peer"
+
+	logging "github.com/ipfs/go-log/v2"
+)
+
+const (
+	EchoService = "test.echo"
+	EchoProtoID = "/test/echo"
+)
+
+var (
+	echoLog = logging.Logger("echo")
+)
+
+type Echo struct {
+	Host host.Host
+
+	WaitBeforeRead, WaitBeforeWrite func() error
+
+	mx     sync.Mutex
+	status EchoStatus
+}
+
+type EchoStatus struct {
+	StreamsIn                 int
+	EchosIn, EchosOut         int
+	IOErrors                  int
+	ResourceServiceErrors     int
+	ResourceReservationErrors int
+}
+
+func NewEcho(h host.Host) *Echo {
+	e := &Echo{Host: h}
+	h.SetStreamHandler(EchoProtoID, e.handleStream)
+	return e
+}
+
+func (e *Echo) Status() EchoStatus {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	return e.status
+}
+
+func (e *Echo) handleStream(s network.Stream) {
+	defer s.Close()
+
+	e.mx.Lock()
+	e.status.StreamsIn++
+	e.mx.Unlock()
+
+	if err := s.Scope().SetService(EchoService); err != nil {
+		echoLog.Debugf("error attaching stream to echo service: %s", err)
+
+		e.mx.Lock()
+		e.status.ResourceServiceErrors++
+		e.mx.Unlock()
+
+		s.Reset()
+		return
+	}
+
+	if err := s.Scope().ReserveMemory(4096, network.ReservationPriorityAlways); err != nil {
+		echoLog.Debugf("error reserving memory: %s", err)
+
+		e.mx.Lock()
+		e.status.ResourceReservationErrors++
+		e.mx.Unlock()
+
+		s.Reset()
+		return
+	}
+
+	if e.WaitBeforeRead != nil {
+		if err := e.WaitBeforeRead(); err != nil {
+			echoLog.Debugf("error waiting before read: %s", err)
+
+			s.Reset()
+			return
+		}
+	}
+
+	buf := make([]byte, 4096)
+
+	s.SetReadDeadline(time.Now().Add(5 * time.Second))
+	n, err := s.Read(buf)
+	switch {
+	case err == io.EOF:
+		if n == 0 {
+			return
+		}
+
+	case err != nil:
+		echoLog.Debugf("I/O error : %s", err)
+
+		e.mx.Lock()
+		e.status.IOErrors++
+		e.mx.Unlock()
+
+		s.Reset()
+		return
+	}
+
+	e.mx.Lock()
+	e.status.EchosIn++
+	e.mx.Unlock()
+
+	if e.WaitBeforeWrite != nil {
+		if err := e.WaitBeforeWrite(); err != nil {
+			echoLog.Debugf("error waiting before write: %s", err)
+
+			s.Reset()
+			return
+		}
+	}
+
+	s.SetWriteDeadline(time.Now().Add(5 * time.Second))
+	_, err = s.Write(buf[:n])
+	if err != nil {
+		echoLog.Debugf("I/O error: %s", err)
+
+		e.mx.Lock()
+		e.status.IOErrors++
+		e.mx.Unlock()
+
+		s.Reset()
+		return
+	}
+
+	e.mx.Lock()
+	e.status.EchosOut++
+	e.mx.Unlock()
+
+	s.CloseWrite()
+}
+
+func (e *Echo) Echo(p peer.ID, what string) error {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	s, err := e.Host.NewStream(ctx, p, EchoProtoID)
+	if err != nil {
+		return err
+	}
+	defer s.Close()
+
+	if err := s.Scope().SetService(EchoService); err != nil {
+		echoLog.Debugf("error attaching stream to echo service: %s", err)
+
+		s.Reset()
+		return err
+	}
+
+	if err := s.Scope().ReserveMemory(4096, network.ReservationPriorityAlways); err != nil {
+		echoLog.Debugf("error reserving memory: %s", err)
+
+		s.Reset()
+		return err
+	}
+
+	s.SetWriteDeadline(time.Now().Add(5 * time.Second))
+	_, err = s.Write([]byte(what))
+	if err != nil {
+		return err
+	}
+	s.CloseWrite()
+
+	buf := make([]byte, 4096)
+
+	s.SetReadDeadline(time.Now().Add(5 * time.Second))
+	n, err := s.Read(buf)
+	switch {
+	case err == io.EOF:
+		if n == 0 {
+			return err
+		}
+
+	case err != nil:
+		echoLog.Debugf("I/O error : %s", err)
+
+		s.Reset()
+		return err
+	}
+
+	if what != string(buf[:n]) {
+		return fmt.Errorf("echo output doesn't match input")
+	}
+
+	return nil
+}

itest/echo_test.go

@@ -0,0 +1,86 @@
+package itest
+
+import (
+	"context"
+	"testing"
+
+	"github.com/libp2p/go-libp2p"
+	"github.com/libp2p/go-libp2p-core/peer"
+	"github.com/libp2p/go-libp2p-core/peerstore"
+)
+
+func createEchos(t *testing.T, count int, opts ...libp2p.Option) []*Echo {
+	result := make([]*Echo, 0, count)
+
+	for i := 0; i < count; i++ {
+		h, err := libp2p.New(opts...)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		e := NewEcho(h)
+		result = append(result, e)
+	}
+
+	for i := 0; i < count; i++ {
+		for j := 0; j < count; j++ {
+			if i == j {
+				continue
+			}
+
+			result[i].Host.Peerstore().AddAddrs(result[j].Host.ID(), result[j].Host.Addrs(), peerstore.PermanentAddrTTL)
+		}
+	}
+
+	return result
+}
+
+func checkEchoStatus(t *testing.T, e *Echo, expected EchoStatus) {
+	t.Helper()
+
+	status := e.Status()
+
+	if status.StreamsIn != expected.StreamsIn {
+		t.Fatalf("expected %d streams in, got %d", expected.StreamsIn, status.StreamsIn)
+	}
+	if status.EchosIn != expected.EchosIn {
+		t.Fatalf("expected %d echos in, got %d", expected.EchosIn, status.EchosIn)
+	}
+	if status.EchosOut != expected.EchosOut {
+		t.Fatalf("expected %d echos out, got %d", expected.EchosOut, status.EchosOut)
+	}
+	if status.IOErrors != expected.IOErrors {
+		t.Fatalf("expected %d I/O errors, got %d", expected.IOErrors, status.IOErrors)
+	}
+	if status.ResourceServiceErrors != expected.ResourceServiceErrors {
+		t.Fatalf("expected %d service resource errors, got %d", expected.ResourceServiceErrors, status.ResourceServiceErrors)
+	}
+	if status.ResourceReservationErrors != expected.ResourceReservationErrors {
+		t.Fatalf("expected %d reservation resource errors, got %d", expected.ResourceReservationErrors, status.ResourceReservationErrors)
+	}
+}
+
+func TestEcho(t *testing.T) {
+	echos := createEchos(t, 2)
+
+	err := echos[0].Host.Connect(context.TODO(), peer.AddrInfo{ID: echos[1].Host.ID()})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		for _, e := range echos {
+			e.Host.Close()
+		}
+	}()
+
+	if err = echos[0].Echo(echos[1].Host.ID(), "hello libp2p"); err != nil {
+		t.Fatal(err)
+	}
+
+	checkEchoStatus(t, echos[1], EchoStatus{
+		StreamsIn: 1,
+		EchosIn:   1,
+		EchosOut:  1,
+	})
+}