testPool: add peers

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
 }