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

     1  package query
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/require"
     7  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
     8  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
     9  	"go.uber.org/mock/gomock"
    10  	grpcCodes "google.golang.org/grpc/codes"
    11  	grpcStatus "google.golang.org/grpc/status"
    12  
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/session"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    17  )
    18  
    19  func TestCreateSession(t *testing.T) {
    20  	trace := &trace.Query{
    21  		OnSessionCreate: func(info trace.QuerySessionCreateStartInfo) func(info trace.QuerySessionCreateDoneInfo) {
    22  			return func(info trace.QuerySessionCreateDoneInfo) {
    23  				if info.Session != nil && info.Error != nil {
    24  					panic("only one result from tuple may be not nil")
    25  				}
    26  			}
    27  		},
    28  	}
    29  	t.Run("HappyWay", func(t *testing.T) {
    30  		ctx := xtest.Context(t)
    31  		ctrl := gomock.NewController(t)
    32  		client := NewMockQueryServiceClient(ctrl)
    33  		client.EXPECT().CreateSession(gomock.Any(), gomock.Any()).Return(&Ydb_Query.CreateSessionResponse{
    34  			Status:    Ydb.StatusIds_SUCCESS,
    35  			SessionId: "123",
    36  		}, nil)
    37  		attachStream := NewMockQueryService_AttachSessionClient(ctrl)
    38  		attachStream.EXPECT().Recv().Return(&Ydb_Query.SessionState{
    39  			Status: Ydb.StatusIds_SUCCESS,
    40  		}, nil).AnyTimes()
    41  		client.EXPECT().AttachSession(gomock.Any(), &Ydb_Query.AttachSessionRequest{
    42  			SessionId: "123",
    43  		}).Return(attachStream, nil)
    44  		require.NotPanics(t, func() {
    45  			s, err := createSession(ctx, client, session.WithTrace(trace))
    46  			require.NoError(t, err)
    47  			require.NotNil(t, s)
    48  			require.Equal(t, "123", s.ID())
    49  		})
    50  	})
    51  	t.Run("TransportError", func(t *testing.T) {
    52  		t.Run("OnCreateSession", func(t *testing.T) {
    53  			ctx := xtest.Context(t)
    54  			ctrl := gomock.NewController(t)
    55  			client := NewMockQueryServiceClient(ctrl)
    56  			client.EXPECT().CreateSession(gomock.Any(), gomock.Any()).Return(nil,
    57  				xerrors.Transport(grpcStatus.Error(grpcCodes.Unavailable, "test")),
    58  			)
    59  			require.NotPanics(t, func() {
    60  				s, err := createSession(ctx, client, session.WithTrace(trace))
    61  				require.Error(t, err)
    62  				require.Nil(t, s)
    63  			})
    64  		})
    65  		t.Run("OnAttachStream", func(t *testing.T) {
    66  			ctx := xtest.Context(t)
    67  			ctrl := gomock.NewController(t)
    68  			client := NewMockQueryServiceClient(ctrl)
    69  			client.EXPECT().CreateSession(gomock.Any(), gomock.Any()).Return(&Ydb_Query.CreateSessionResponse{
    70  				Status:    Ydb.StatusIds_SUCCESS,
    71  				SessionId: "123",
    72  			}, nil)
    73  			client.EXPECT().AttachSession(gomock.Any(), &Ydb_Query.AttachSessionRequest{
    74  				SessionId: "123",
    75  			}).Return(nil, xerrors.Transport(grpcStatus.Error(grpcCodes.Unavailable, "test")))
    76  			client.EXPECT().DeleteSession(gomock.Any(), &Ydb_Query.DeleteSessionRequest{
    77  				SessionId: "123",
    78  			}).Return(&Ydb_Query.DeleteSessionResponse{
    79  				Status: Ydb.StatusIds_SUCCESS,
    80  			}, nil)
    81  			require.NotPanics(t, func() {
    82  				s, err := createSession(ctx, client, session.WithTrace(trace))
    83  				require.Error(t, err)
    84  				require.Nil(t, s)
    85  			})
    86  		})
    87  	})
    88  	t.Run("OperationError", func(t *testing.T) {
    89  		t.Run("OnCreateSession", func(t *testing.T) {
    90  			ctx := xtest.Context(t)
    91  			ctrl := gomock.NewController(t)
    92  			client := NewMockQueryServiceClient(ctrl)
    93  			client.EXPECT().CreateSession(gomock.Any(), gomock.Any()).Return(nil,
    94  				xerrors.Operation(xerrors.WithStatusCode(Ydb.StatusIds_UNAVAILABLE)),
    95  			)
    96  			require.NotPanics(t, func() {
    97  				s, err := createSession(ctx, client, session.WithTrace(trace))
    98  				require.Error(t, err)
    99  				require.Nil(t, s)
   100  			})
   101  		})
   102  		t.Run("OnAttachStream", func(t *testing.T) {
   103  			ctx := xtest.Context(t)
   104  			ctrl := gomock.NewController(t)
   105  			client := NewMockQueryServiceClient(ctrl)
   106  			client.EXPECT().CreateSession(gomock.Any(), gomock.Any()).Return(&Ydb_Query.CreateSessionResponse{
   107  				Status:    Ydb.StatusIds_SUCCESS,
   108  				SessionId: "123",
   109  			}, nil)
   110  			client.EXPECT().AttachSession(gomock.Any(), &Ydb_Query.AttachSessionRequest{
   111  				SessionId: "123",
   112  			}).Return(nil, xerrors.Operation(xerrors.WithStatusCode(Ydb.StatusIds_UNAVAILABLE)))
   113  			client.EXPECT().DeleteSession(gomock.Any(), &Ydb_Query.DeleteSessionRequest{
   114  				SessionId: "123",
   115  			}).Return(&Ydb_Query.DeleteSessionResponse{
   116  				Status: Ydb.StatusIds_SUCCESS,
   117  			}, nil)
   118  			require.NotPanics(t, func() {
   119  				s, err := createSession(ctx, client, session.WithTrace(trace))
   120  				require.Error(t, err)
   121  				require.Nil(t, s)
   122  			})
   123  		})
   124  	})
   125  }