Add WithServer for TestPool

net/rpc/rpctest/pool.go

@@ -16,6 +16,7 @@ func NewTestPool() *TestPool {
 type TestPool struct {
 	peers map[string]peer.Peer
 	mu    sync.Mutex
+	ts    *TestServer
 }
 
 func (t *TestPool) Init(a *app.App) (err error) {
@@ -26,6 +27,13 @@ func (t *TestPool) Name() (name string) {
 	return pool.CName
 }
 
+func (t *TestPool) WithServer(ts *TestServer) *TestPool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	t.ts = ts
+	return t
+}
+
 func (t *TestPool) Run(ctx context.Context) (err error) {
 	return nil
 }
@@ -40,16 +48,24 @@ func (t *TestPool) Get(ctx context.Context, id string) (peer.Peer, error) {
 	if p, ok := t.peers[id]; ok {
 		return p, nil
 	}
-	return nil, net.ErrUnableToConnect
+	if t.ts == nil {
+		return nil, net.ErrUnableToConnect
+	}
+	return t.ts.Dial(id)
 }
 
 func (t *TestPool) GetOneOf(ctx context.Context, peerIds []string) (peer.Peer, error) {
-	for _, id := range peerIds {
+	t.mu.Lock()
-		if p, err := t.Get(ctx, id); err == nil {
+	defer t.mu.Unlock()
+	for _, peerId := range peerIds {
+		if p, ok := t.peers[peerId]; ok {
 			return p, nil
 		}
 	}
-	return nil, net.ErrUnableToConnect
+	if t.ts == nil || len(peerIds) == 0 {
+		return nil, net.ErrUnableToConnect
+	}
+	return t.ts.Dial(peerIds[0])
 }
 
 func (t *TestPool) AddPeer(ctx context.Context, p peer.Peer) (err error) {

net/rpc/rpctest/server.go

@@ -3,6 +3,7 @@ package rpctest
 import (
 	"context"
 	"github.com/anyproto/any-sync/app"
+	"github.com/anyproto/any-sync/net/peer"
 	"github.com/anyproto/any-sync/net/rpc"
 	"github.com/anyproto/any-sync/net/rpc/server"
 	"net"
@@ -10,6 +11,17 @@ import (
 	"storj.io/drpc/drpcserver"
 )
 
+type mockCtrl struct {
+}
+
+func (m mockCtrl) ServeConn(ctx context.Context, conn net.Conn) (err error) {
+	return nil
+}
+
+func (m mockCtrl) DrpcConfig() rpc.Config {
+	return rpc.Config{}
+}
+
 func NewTestServer() *TestServer {
 	ts := &TestServer{
 		Mux: drpcmux.New(),
@@ -23,19 +35,19 @@ type TestServer struct {
 	*drpcserver.Server
 }
 
-func (ts *TestServer) Init(a *app.App) (err error) {
+func (s *TestServer) Init(a *app.App) (err error) {
 	return nil
 }
 
-func (ts *TestServer) Name() (name string) {
+func (s *TestServer) Name() (name string) {
 	return server.CName
 }
 
-func (ts *TestServer) Run(ctx context.Context) (err error) {
+func (s *TestServer) Run(ctx context.Context) (err error) {
 	return nil
 }
 
-func (ts *TestServer) Close(ctx context.Context) (err error) {
+func (s *TestServer) Close(ctx context.Context) (err error) {
 	return nil
 }
 
@@ -46,3 +58,14 @@ func (s *TestServer) ServeConn(ctx context.Context, conn net.Conn) (err error) {
 func (s *TestServer) DrpcConfig() rpc.Config {
 	return rpc.Config{Stream: rpc.StreamConfig{MaxMsgSizeMb: 10}}
 }
+
+func (s *TestServer) Dial(peerId string) (clientPeer peer.Peer, err error) {
+	mcS, mcC := MultiConnPair(peerId+"server", peerId)
+	// NewPeer runs the accept loop
+	_, err = peer.NewPeer(mcS, s)
+	if err != nil {
+		return
+	}
+	// and we ourselves don't call server methods on accept
+	return peer.NewPeer(mcC, mockCtrl{})
+}