basic rcmgr integration tests

itest/echo.go

@@ -26,10 +26,10 @@ var (
 type Echo struct {
 	Host host.Host
 
-	WaitBeforeRead, WaitBeforeWrite func() error
-
 	mx     sync.Mutex
 	status EchoStatus
+
+	beforeReserve, beforeRead, beforeWrite, beforeDone func() error
 }
 
 type EchoStatus struct {
@@ -53,13 +53,78 @@ func (e *Echo) Status() EchoStatus {
 	return e.status
 }
 
+func (e *Echo) BeforeReserve(f func() error) {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	e.beforeReserve = f
+}
+
+func (e *Echo) BeforeRead(f func() error) {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	e.beforeRead = f
+}
+
+func (e *Echo) BeforeWrite(f func() error) {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	e.beforeWrite = f
+}
+
+func (e *Echo) BeforeDone(f func() error) {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	e.beforeDone = f
+}
+
+func (e *Echo) getBeforeReserve() func() error {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	return e.beforeReserve
+}
+
+func (e *Echo) getBeforeRead() func() error {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	return e.beforeRead
+}
+
+func (e *Echo) getBeforeWrite() func() error {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	return e.beforeWrite
+}
+
+func (e *Echo) getBeforeDone() func() error {
+	e.mx.Lock()
+	defer e.mx.Unlock()
+
+	return e.beforeDone
+}
+
 func (e *Echo) handleStream(s network.Stream) {
 	defer s.Close()
 
 	e.mx.Lock()
 	e.status.StreamsIn++
 	e.mx.Unlock()
 
+	if beforeReserve := e.getBeforeReserve(); beforeReserve != nil {
+		if err := beforeReserve(); err != nil {
+			echoLog.Debugf("error syncing before reserve: %s", err)
+
+			s.Reset()
+			return
+		}
+	}
+
 	if err := s.Scope().SetService(EchoService); err != nil {
 		echoLog.Debugf("error attaching stream to echo service: %s", err)
 
@@ -82,9 +147,9 @@ func (e *Echo) handleStream(s network.Stream) {
 		return
 	}
 
-	if e.WaitBeforeRead != nil {
-		if err := e.WaitBeforeRead(); err != nil {
-			echoLog.Debugf("error waiting before read: %s", err)
+	if beforeRead := e.getBeforeRead(); beforeRead != nil {
+		if err := beforeRead(); err != nil {
+			echoLog.Debugf("error syncing before read: %s", err)
 
 			s.Reset()
 			return
@@ -116,9 +181,9 @@ func (e *Echo) handleStream(s network.Stream) {
 	e.status.EchosIn++
 	e.mx.Unlock()
 
-	if e.WaitBeforeWrite != nil {
-		if err := e.WaitBeforeWrite(); err != nil {
-			echoLog.Debugf("error waiting before write: %s", err)
+	if beforeWrite := e.getBeforeWrite(); beforeWrite != nil {
+		if err := beforeWrite(); err != nil {
+			echoLog.Debugf("error syncing before write: %s", err)
 
 			s.Reset()
 			return
@@ -143,6 +208,14 @@ func (e *Echo) handleStream(s network.Stream) {
 	e.mx.Unlock()
 
 	s.CloseWrite()
+
+	if beforeDone := e.getBeforeDone(); beforeDone != nil {
+		if err := beforeDone(); err != nil {
+			echoLog.Debugf("error syncing before done: %s", err)
+
+			s.Reset()
+		}
+	}
 }
 
 func (e *Echo) Echo(p peer.ID, what string) error {

itest/echo_test.go

@@ -7,12 +7,19 @@ import (
 	"github.com/libp2p/go-libp2p"
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/peerstore"
+
+	"github.com/stretchr/testify/require"
 )
 
-func createEchos(t *testing.T, count int, opts ...libp2p.Option) []*Echo {
+func createEchos(t *testing.T, count int, makeOpts ...func(int) libp2p.Option) []*Echo {
 	result := make([]*Echo, 0, count)
 
 	for i := 0; i < count; i++ {
+		opts := make([]libp2p.Option, 0, len(makeOpts))
+		for _, makeOpt := range makeOpts {
+			opts = append(opts, makeOpt(i))
+		}
+
 		h, err := libp2p.New(opts...)
 		if err != nil {
 			t.Fatal(err)
@@ -35,46 +42,26 @@ func createEchos(t *testing.T, count int, opts ...libp2p.Option) []*Echo {
 	return result
 }
 
+func closeEchos(echos []*Echo) {
+	for _, e := range echos {
+		e.Host.Close()
+	}
+}
+
 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)
-	}
+	require.Equal(t, expected, e.Status())
 }
 
 func TestEcho(t *testing.T) {
 	echos := createEchos(t, 2)
+	defer closeEchos(echos)
 
-	err := echos[0].Host.Connect(context.TODO(), peer.AddrInfo{ID: echos[1].Host.ID()})
-	if err != nil {
+	if err := echos[0].Host.Connect(context.TODO(), peer.AddrInfo{ID: echos[1].Host.ID()}); 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 {
+	if err := echos[0].Echo(echos[1].Host.ID(), "hello libp2p"); err != nil {
 		t.Fatal(err)
 	}