From 35c29a48421434a121f86bd1127a7786c204df6f Mon Sep 17 00:00:00 2001
From: mcrakhman <rakhmanov.m@gmail.com>
Date: Thu, 8 Jun 2023 21:07:03 +0200
Subject: [PATCH] Add WithServer for TestPool

---
 net/rpc/rpctest/pool.go   | 24 ++++++++++++++++++++----
 net/rpc/rpctest/server.go | 31 +++++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/net/rpc/rpctest/pool.go b/net/rpc/rpctest/pool.go
index 0e2b5d1b..fc70d16b 100644
--- a/net/rpc/rpctest/pool.go
+++ b/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 {
-		if p, err := t.Get(ctx, id); err == nil {
+	t.mu.Lock()
+	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) {
diff --git a/net/rpc/rpctest/server.go b/net/rpc/rpctest/server.go
index 5c03f3e4..de6bf68b 100644
--- a/net/rpc/rpctest/server.go
+++ b/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{})
+}