From 2b8329e026c74d310817938c6cb8243749d0d49e Mon Sep 17 00:00:00 2001
From: Alex Vanin <a.vanin@yadro.com>
Date: Fri, 28 Feb 2025 12:29:23 +0300
Subject: [PATCH] [#336] pool/tree: Increase test coverage in TestStreamRetry

Signed-off-by: Alex Vanin <a.vanin@yadro.com>
---
 pool/tree/pool_server_test.go | 99 ++++++++++++++++++++++++++++++-----
 1 file changed, 87 insertions(+), 12 deletions(-)

diff --git a/pool/tree/pool_server_test.go b/pool/tree/pool_server_test.go
index 014e454..01373da 100644
--- a/pool/tree/pool_server_test.go
+++ b/pool/tree/pool_server_test.go
@@ -4,12 +4,14 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"io"
 	"net"
 	"runtime"
 	"strconv"
 	"testing"
 
 	apinetmap "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/netmap"
+	apitree "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/tree"
 	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
 	cidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id/test"
 	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/netmap"
@@ -29,8 +31,9 @@ type mockTreeServer struct {
 	healthy    bool
 	addCounter int
 
-	getSubTreeError   error
-	getSubTreeCounter int
+	getSubTreeError     error
+	getSubTreeResponses []*tree.GetSubTreeResponse_Body
+	getSubTreeCounter   int
 }
 
 type mockNetmapSource struct {
@@ -93,9 +96,22 @@ func (m *mockTreeServer) GetNodeByPath(context.Context, *tree.GetNodeByPathReque
 	panic("implement me")
 }
 
-func (m *mockTreeServer) GetSubTree(*tree.GetSubTreeRequest, tree.TreeService_GetSubTreeServer) error {
+func (m *mockTreeServer) GetSubTree(_ *tree.GetSubTreeRequest, s tree.TreeService_GetSubTreeServer) error {
 	m.getSubTreeCounter++
-	return m.getSubTreeError
+
+	if m.getSubTreeError != nil {
+		return m.getSubTreeError
+	}
+
+	for i := range m.getSubTreeResponses {
+		if err := s.Send(&tree.GetSubTreeResponse{
+			Body: m.getSubTreeResponses[i],
+		}); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
 func (m *mockTreeServer) TreeList(context.Context, *tree.TreeListRequest) (*tree.TreeListResponse, error) {
@@ -246,11 +262,20 @@ func TestStreamRetry(t *testing.T) {
 		placementPolicy = "REP 2"
 	)
 
-	// Initialize gRPC servers and create pool with netmap source
-	treePool, servers, _ := preparePoolWithNetmapSource(t, numberOfNodes, placementPolicy)
-	for i := range servers {
-		servers[i].getSubTreeError = errors.New("tree not found")
+	expected := []*tree.GetSubTreeResponse_Body{
+		{
+			NodeId: []uint64{1},
+		},
+		{
+			NodeId: []uint64{2},
+		},
+		{
+			NodeId: []uint64{3},
+		},
 	}
+
+	// Initialize gRPC servers and create pool with netmap source
+	treePool, servers, source := preparePoolWithNetmapSource(t, numberOfNodes, placementPolicy)
 	defer func() {
 		for i := range servers {
 			servers[i].Stop()
@@ -260,11 +285,61 @@ func TestStreamRetry(t *testing.T) {
 	cnr := cidtest.ID()
 	ctx := context.Background()
 
-	_, err := treePool.GetSubTree(ctx, GetSubTreeParams{CID: cnr})
-	require.Error(t, err)
+	sortedServers, err := sortServers(ctx, servers, source, cnr)
+	require.NoError(t, err)
+
+	// Return expected response in last priority node, others return error
+	for i := range sortedServers {
+		if i == len(sortedServers)-1 {
+			sortedServers[i].getSubTreeResponses = expected
+		} else {
+			sortedServers[i].getSubTreeError = errors.New("tree not found")
+		}
+	}
+
+	t.Run("read all", func(t *testing.T) {
+		reader, err := treePool.GetSubTree(ctx, GetSubTreeParams{CID: cnr})
+		require.NoError(t, err)
+
+		data, err := reader.ReadAll()
+		require.NoError(t, err)
+
+		require.Len(t, data, len(expected))
+		for i := range expected {
+			require.EqualValues(t, expected[i].GetNodeId(), data[i].GetNodeID())
+		}
+	})
+
+	t.Run("next", func(t *testing.T) {
+		reader, err := treePool.GetSubTree(ctx, GetSubTreeParams{CID: cnr})
+		require.NoError(t, err)
+
+		for i := range expected {
+			resp, err := reader.Next()
+			require.NoError(t, err)
+			require.Equal(t, expected[i].GetNodeId(), resp.GetNodeID())
+		}
+
+		_, err = reader.Next()
+		require.Error(t, io.EOF, err)
+	})
+
+	t.Run("read", func(t *testing.T) {
+		reader, err := treePool.GetSubTree(ctx, GetSubTreeParams{CID: cnr})
+		require.NoError(t, err)
+
+		buf := make([]*apitree.GetSubTreeResponseBody, len(expected))
+		_, err = reader.Read(buf)
+		require.NoError(t, err)
+
+		require.Len(t, buf, len(expected))
+		for i := range expected {
+			require.EqualValues(t, expected[i].GetNodeId(), buf[i].GetNodeID())
+		}
+	})
 
 	for i := range servers {
-		// check we retried every available node in the pool
-		require.Equal(t, 1, servers[i].getSubTreeCounter)
+		// check we retried every available node in the pool three times
+		require.Equal(t, 3, servers[i].getSubTreeCounter)
 	}
 }