github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/query/result_go1.23_test.go (about)

     1  //go:build go1.23
     2  
     3  package query
     4  
     5  import (
     6  	"context"
     7  	"io"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    12  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
    13  	"go.uber.org/mock/gomock"
    14  
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    16  )
    17  
    18  func TestResultRangeResultSets(t *testing.T) {
    19  	ctx, cancel := context.WithCancel(xtest.Context(t))
    20  	defer cancel()
    21  	ctrl := gomock.NewController(t)
    22  	stream := NewMockQueryService_ExecuteQueryClient(ctrl)
    23  	stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    24  		Status:         Ydb.StatusIds_SUCCESS,
    25  		ResultSetIndex: 0,
    26  		ResultSet: &Ydb.ResultSet{
    27  			Columns: []*Ydb.Column{
    28  				{
    29  					Name: "a",
    30  					Type: &Ydb.Type{
    31  						Type: &Ydb.Type_TypeId{
    32  							TypeId: Ydb.Type_UINT64,
    33  						},
    34  					},
    35  				},
    36  				{
    37  					Name: "b",
    38  					Type: &Ydb.Type{
    39  						Type: &Ydb.Type_TypeId{
    40  							TypeId: Ydb.Type_UTF8,
    41  						},
    42  					},
    43  				},
    44  			},
    45  			Rows: []*Ydb.Value{
    46  				{
    47  					Items: []*Ydb.Value{{
    48  						Value: &Ydb.Value_Uint64Value{
    49  							Uint64Value: 1,
    50  						},
    51  					}, {
    52  						Value: &Ydb.Value_TextValue{
    53  							TextValue: "1",
    54  						},
    55  					}},
    56  				},
    57  				{
    58  					Items: []*Ydb.Value{{
    59  						Value: &Ydb.Value_Uint64Value{
    60  							Uint64Value: 2,
    61  						},
    62  					}, {
    63  						Value: &Ydb.Value_TextValue{
    64  							TextValue: "2",
    65  						},
    66  					}},
    67  				},
    68  				{
    69  					Items: []*Ydb.Value{{
    70  						Value: &Ydb.Value_Uint64Value{
    71  							Uint64Value: 3,
    72  						},
    73  					}, {
    74  						Value: &Ydb.Value_TextValue{
    75  							TextValue: "3",
    76  						},
    77  					}},
    78  				},
    79  			},
    80  		},
    81  	}, nil)
    82  	stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
    83  		Status:         Ydb.StatusIds_SUCCESS,
    84  		ResultSetIndex: 0,
    85  		ResultSet: &Ydb.ResultSet{
    86  			Rows: []*Ydb.Value{
    87  				{
    88  					Items: []*Ydb.Value{{
    89  						Value: &Ydb.Value_Uint64Value{
    90  							Uint64Value: 4,
    91  						},
    92  					}, {
    93  						Value: &Ydb.Value_TextValue{
    94  							TextValue: "4",
    95  						},
    96  					}},
    97  				},
    98  				{
    99  					Items: []*Ydb.Value{{
   100  						Value: &Ydb.Value_Uint64Value{
   101  							Uint64Value: 5,
   102  						},
   103  					}, {
   104  						Value: &Ydb.Value_TextValue{
   105  							TextValue: "5",
   106  						},
   107  					}},
   108  				},
   109  			},
   110  		},
   111  	}, nil)
   112  	stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   113  		Status:         Ydb.StatusIds_SUCCESS,
   114  		ResultSetIndex: 1,
   115  		ResultSet: &Ydb.ResultSet{
   116  			Columns: []*Ydb.Column{
   117  				{
   118  					Name: "c",
   119  					Type: &Ydb.Type{
   120  						Type: &Ydb.Type_TypeId{
   121  							TypeId: Ydb.Type_UINT64,
   122  						},
   123  					},
   124  				},
   125  				{
   126  					Name: "d",
   127  					Type: &Ydb.Type{
   128  						Type: &Ydb.Type_TypeId{
   129  							TypeId: Ydb.Type_UTF8,
   130  						},
   131  					},
   132  				},
   133  				{
   134  					Name: "e",
   135  					Type: &Ydb.Type{
   136  						Type: &Ydb.Type_TypeId{
   137  							TypeId: Ydb.Type_BOOL,
   138  						},
   139  					},
   140  				},
   141  			},
   142  			Rows: []*Ydb.Value{
   143  				{
   144  					Items: []*Ydb.Value{{
   145  						Value: &Ydb.Value_Uint64Value{
   146  							Uint64Value: 1,
   147  						},
   148  					}, {
   149  						Value: &Ydb.Value_TextValue{
   150  							TextValue: "1",
   151  						},
   152  					}, {
   153  						Value: &Ydb.Value_BoolValue{
   154  							BoolValue: true,
   155  						},
   156  					}},
   157  				},
   158  				{
   159  					Items: []*Ydb.Value{{
   160  						Value: &Ydb.Value_Uint64Value{
   161  							Uint64Value: 2,
   162  						},
   163  					}, {
   164  						Value: &Ydb.Value_TextValue{
   165  							TextValue: "2",
   166  						},
   167  					}, {
   168  						Value: &Ydb.Value_BoolValue{
   169  							BoolValue: false,
   170  						},
   171  					}},
   172  				},
   173  			},
   174  		},
   175  	}, nil)
   176  	stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   177  		Status:         Ydb.StatusIds_SUCCESS,
   178  		ResultSetIndex: 1,
   179  		ResultSet: &Ydb.ResultSet{
   180  			Rows: []*Ydb.Value{
   181  				{
   182  					Items: []*Ydb.Value{{
   183  						Value: &Ydb.Value_Uint64Value{
   184  							Uint64Value: 3,
   185  						},
   186  					}, {
   187  						Value: &Ydb.Value_TextValue{
   188  							TextValue: "3",
   189  						},
   190  					}, {
   191  						Value: &Ydb.Value_BoolValue{
   192  							BoolValue: true,
   193  						},
   194  					}},
   195  				},
   196  				{
   197  					Items: []*Ydb.Value{{
   198  						Value: &Ydb.Value_Uint64Value{
   199  							Uint64Value: 4,
   200  						},
   201  					}, {
   202  						Value: &Ydb.Value_TextValue{
   203  							TextValue: "4",
   204  						},
   205  					}, {
   206  						Value: &Ydb.Value_BoolValue{
   207  							BoolValue: false,
   208  						},
   209  					}},
   210  				},
   211  				{
   212  					Items: []*Ydb.Value{{
   213  						Value: &Ydb.Value_Uint64Value{
   214  							Uint64Value: 5,
   215  						},
   216  					}, {
   217  						Value: &Ydb.Value_TextValue{
   218  							TextValue: "5",
   219  						},
   220  					}, {
   221  						Value: &Ydb.Value_BoolValue{
   222  							BoolValue: false,
   223  						},
   224  					}},
   225  				},
   226  			},
   227  		},
   228  	}, nil)
   229  	stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   230  		Status:         Ydb.StatusIds_SUCCESS,
   231  		ResultSetIndex: 2,
   232  		ResultSet: &Ydb.ResultSet{
   233  			Columns: []*Ydb.Column{
   234  				{
   235  					Name: "c",
   236  					Type: &Ydb.Type{
   237  						Type: &Ydb.Type_TypeId{
   238  							TypeId: Ydb.Type_UINT64,
   239  						},
   240  					},
   241  				},
   242  				{
   243  					Name: "d",
   244  					Type: &Ydb.Type{
   245  						Type: &Ydb.Type_TypeId{
   246  							TypeId: Ydb.Type_UTF8,
   247  						},
   248  					},
   249  				},
   250  				{
   251  					Name: "e",
   252  					Type: &Ydb.Type{
   253  						Type: &Ydb.Type_TypeId{
   254  							TypeId: Ydb.Type_BOOL,
   255  						},
   256  					},
   257  				},
   258  			},
   259  			Rows: []*Ydb.Value{
   260  				{
   261  					Items: []*Ydb.Value{{
   262  						Value: &Ydb.Value_Uint64Value{
   263  							Uint64Value: 1,
   264  						},
   265  					}, {
   266  						Value: &Ydb.Value_TextValue{
   267  							TextValue: "1",
   268  						},
   269  					}, {
   270  						Value: &Ydb.Value_BoolValue{
   271  							BoolValue: true,
   272  						},
   273  					}},
   274  				},
   275  				{
   276  					Items: []*Ydb.Value{{
   277  						Value: &Ydb.Value_Uint64Value{
   278  							Uint64Value: 2,
   279  						},
   280  					}, {
   281  						Value: &Ydb.Value_TextValue{
   282  							TextValue: "2",
   283  						},
   284  					}, {
   285  						Value: &Ydb.Value_BoolValue{
   286  							BoolValue: false,
   287  						},
   288  					}},
   289  				},
   290  			},
   291  		},
   292  	}, nil)
   293  	stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
   294  		Status:         Ydb.StatusIds_SUCCESS,
   295  		ResultSetIndex: 2,
   296  		ResultSet: &Ydb.ResultSet{
   297  			Rows: []*Ydb.Value{
   298  				{
   299  					Items: []*Ydb.Value{{
   300  						Value: &Ydb.Value_Uint64Value{
   301  							Uint64Value: 3,
   302  						},
   303  					}, {
   304  						Value: &Ydb.Value_TextValue{
   305  							TextValue: "3",
   306  						},
   307  					}, {
   308  						Value: &Ydb.Value_BoolValue{
   309  							BoolValue: true,
   310  						},
   311  					}},
   312  				},
   313  				{
   314  					Items: []*Ydb.Value{{
   315  						Value: &Ydb.Value_Uint64Value{
   316  							Uint64Value: 4,
   317  						},
   318  					}, {
   319  						Value: &Ydb.Value_TextValue{
   320  							TextValue: "4",
   321  						},
   322  					}, {
   323  						Value: &Ydb.Value_BoolValue{
   324  							BoolValue: false,
   325  						},
   326  					}},
   327  				},
   328  				{
   329  					Items: []*Ydb.Value{{
   330  						Value: &Ydb.Value_Uint64Value{
   331  							Uint64Value: 5,
   332  						},
   333  					}, {
   334  						Value: &Ydb.Value_TextValue{
   335  							TextValue: "5",
   336  						},
   337  					}, {
   338  						Value: &Ydb.Value_BoolValue{
   339  							BoolValue: false,
   340  						},
   341  					}},
   342  				},
   343  			},
   344  		},
   345  	}, nil)
   346  	stream.EXPECT().Recv().Return(nil, io.EOF)
   347  	r, err := newResult(ctx, stream, nil)
   348  	require.NoError(t, err)
   349  	defer r.Close(ctx)
   350  	rsCount := 0
   351  	for rs, err := range r.ResultSets(ctx) {
   352  		require.NoError(t, err)
   353  		rowsCount := 0
   354  		for _, err := range rs.Rows(ctx) {
   355  			require.NoError(t, err)
   356  			rowsCount++
   357  		}
   358  		require.EqualValues(t, 5, rowsCount)
   359  		rsCount++
   360  	}
   361  	require.EqualValues(t, 3, rsCount)
   362  }