[#114] tree: Fix retry tests

Signed-off-by: Denis Kirillov <d.kirillov@yadro.com>

internal/frostfs/services/tree_client_grpc.go

@@ -68,10 +68,15 @@ func (n GetSubTreeResponseBodyWrapper) GetMeta() []tree.Meta {
 	return res
 }
 
+type TreeClient interface {
+	TreeClient(ctx context.Context) (grpcService.TreeServiceClient, error)
+	Address() string
+}
+
 type ServiceClientGRPC struct {
 	key        *keys.PrivateKey
 	log        *zap.Logger
-	clients    []*treeClient.TreeClient
+	clients    []TreeClient
 	startIndex int32
 }
 
@@ -99,7 +104,7 @@ func NewTreeServiceClientGRPC(ctx context.Context, endpoints []string, key *keys
 
 	firstHealthy := -1
 
-	res.clients = make([]*treeClient.TreeClient, len(endpoints))
+	res.clients = make([]TreeClient, len(endpoints))
 	for i, addr := range endpoints {
 		res.clients[i] = treeClient.NewTreeClient(addr, grpcOpts...)
 		if _, err := res.clients[i].TreeClient(ctx); err != nil {

internal/frostfs/services/tree_client_grpc_test.go

@@ -3,17 +3,30 @@ package services
 import (
 	"context"
 	"errors"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
 	"testing"
 
 	"git.frostfs.info/TrueCloudLab/frostfs-s3-gw/api/layer"
-	"git.frostfs.info/TrueCloudLab/frostfs-s3-gw/internal/frostfs/services/client"
 	grpcService "git.frostfs.info/TrueCloudLab/frostfs-s3-gw/internal/frostfs/services/tree"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap/zaptest"
 )
 
+type treeClientMock struct {
+	address string
+	err     bool
+}
+
+func (t *treeClientMock) TreeClient(context.Context) (grpcService.TreeServiceClient, error) {
+	if t.err {
+		return nil, errors.New("error")
+	}
+	return nil, nil
+}
+
+func (t *treeClientMock) Address() string {
+	return t.address
+}
+
 func TestHandleError(t *testing.T) {
 	defaultError := errors.New("default error")
 	for _, tc := range []struct {
@@ -44,58 +57,73 @@ func TestRetry(t *testing.T) {
 	ctx := context.Background()
 	log := zaptest.NewLogger(t)
 
-	grpcDialOpt := grpc.WithTransportCredentials(insecure.NewCredentials())
 	cl := &ServiceClientGRPC{
 		log: zaptest.NewLogger(t),
-		clients: []*client.TreeClient{
+		clients: []TreeClient{
-			client.NewTreeClient("node0", grpcDialOpt),
+			&treeClientMock{address: "node0"},
-			client.NewTreeClient("node1", grpcDialOpt),
+			&treeClientMock{address: "node1"},
-			client.NewTreeClient("node2", grpcDialOpt),
+			&treeClientMock{address: "node2"},
-			client.NewTreeClient("node3", grpcDialOpt),
+			&treeClientMock{address: "node3"},
 		},
 	}
 
-	makeFn := func(shouldFail []string) func(grpcService.TreeServiceClient) error {
+	makeFn := func(client grpcService.TreeServiceClient) error {
-		return func(client grpcService.TreeServiceClient) error {
+		return nil
-
-			//for _, item := range shouldFail {
-			//if item == client.address {
-			//	return errors.New("not found")
-			//}
-			//}
-			return nil
-		}
 	}
 
 	t.Run("first ok", func(t *testing.T) {
-		err := cl.requestWithRetry(ctx, log, makeFn([]string{}))
+		err := cl.requestWithRetry(ctx, log, makeFn)
 		require.NoError(t, err)
 		require.Equal(t, 0, cl.getStartIndex())
-		cl.setStartIndex(0)
+		resetClients(cl)
 	})
 
 	t.Run("first failed", func(t *testing.T) {
-		err := cl.requestWithRetry(ctx, log, makeFn([]string{"node0"}))
+		setErrors(cl.clients[:1])
+		err := cl.requestWithRetry(ctx, log, makeFn)
 		require.NoError(t, err)
 		require.Equal(t, 1, cl.getStartIndex())
-		cl.setStartIndex(0)
+		resetClients(cl)
 	})
 
 	t.Run("all failed", func(t *testing.T) {
-		err := cl.requestWithRetry(ctx, log, makeFn([]string{"node0", "node1", "node2", "node3"}))
+		setErrors(cl.clients)
+		err := cl.requestWithRetry(ctx, log, makeFn)
 		require.Error(t, err)
 		require.Equal(t, 0, cl.getStartIndex())
-		cl.setStartIndex(0)
+		resetClients(cl)
 	})
 
 	t.Run("round", func(t *testing.T) {
-		err := cl.requestWithRetry(ctx, log, makeFn([]string{"node0", "node1"}))
+		setErrors(cl.clients[:2])
+		err := cl.requestWithRetry(ctx, log, makeFn)
 		require.NoError(t, err)
 		require.Equal(t, 2, cl.getStartIndex())
+		resetClientsErrors(cl)
 
-		err = cl.requestWithRetry(ctx, log, makeFn([]string{"node2", "node3"}))
+		setErrors(cl.clients[2:])
+		err = cl.requestWithRetry(ctx, log, makeFn)
 		require.NoError(t, err)
 		require.Equal(t, 0, cl.getStartIndex())
-		cl.setStartIndex(0)
+		resetClients(cl)
 	})
 }
+
+func resetClients(cl *ServiceClientGRPC) {
+	resetClientsErrors(cl)
+	cl.setStartIndex(0)
+}
+
+func resetClientsErrors(cl *ServiceClientGRPC) {
+	for _, client := range cl.clients {
+		node := client.(*treeClientMock)
+		node.err = false
+	}
+}
+
+func setErrors(clients []TreeClient) {
+	for _, client := range clients {
+		node := client.(*treeClientMock)
+		node.err = true
+	}
+}