From 6d2178dfc95d29908dc31bf0932ce7841bc568cf Mon Sep 17 00:00:00 2001
From: Sergey Cherepanov <cheggaaa@gmail.com>
Date: Thu, 2 Feb 2023 13:18:30 +0300
Subject: [PATCH] testPool: add peers

---
 net/rpc/rpctest/pool.go | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/net/rpc/rpctest/pool.go b/net/rpc/rpctest/pool.go
index e7d0893d..5faedd44 100644
--- a/net/rpc/rpctest/pool.go
+++ b/net/rpc/rpctest/pool.go
@@ -15,12 +15,15 @@ import (
 var ErrCantConnect = errors.New("can't connect to test server")
 
 func NewTestPool() *TestPool {
-	return &TestPool{}
+	return &TestPool{
+		peers: map[string]peer.Peer{},
+	}
 }
 
 type TestPool struct {
-	ts *TesServer
-	mu sync.Mutex
+	ts    *TesServer
+	peers map[string]peer.Peer
+	mu    sync.Mutex
 }
 
 func (t *TestPool) WithServer(ts *TesServer) *TestPool {
@@ -33,6 +36,9 @@ func (t *TestPool) WithServer(ts *TesServer) *TestPool {
 func (t *TestPool) Get(ctx context.Context, id string) (peer.Peer, error) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
+	if p, ok := t.peers[id]; ok {
+		return p, nil
+	}
 	if t.ts == nil {
 		return nil, ErrCantConnect
 	}
@@ -51,6 +57,11 @@ func (t *TestPool) Dial(ctx context.Context, id string) (peer.Peer, error) {
 func (t *TestPool) GetOneOf(ctx context.Context, peerIds []string) (peer.Peer, error) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
+	for _, peerId := range peerIds {
+		if p, ok := t.peers[peerId]; ok {
+			return p, nil
+		}
+	}
 	if t.ts == nil {
 		return nil, ErrCantConnect
 	}
@@ -70,6 +81,12 @@ func (t *TestPool) NewPool(name string) pool.Pool {
 	return t
 }
 
+func (t *TestPool) AddPeer(p peer.Peer) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	t.peers[p.Id()] = p
+}
+
 func (t *TestPool) Init(a *app.App) (err error) {
 	return nil
 }