@@ -10,47 +10,162 @@ import (
1010 grpcCodes "google.golang.org/grpc/codes"
1111 grpcStatus "google.golang.org/grpc/status"
1212
13+ "github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
1314 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
1415 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
1516 "github.com/ydb-platform/ydb-go-sdk/v3/query"
17+ "github.com/ydb-platform/ydb-go-sdk/v3/trace"
1618)
1719
1820func TestBegin (t * testing.T ) {
1921 t .Run ("HappyWay" , func (t * testing.T ) {
2022 ctx := xtest .Context (t )
2123 ctrl := gomock .NewController (t )
22- service := NewMockQueryServiceClient (ctrl )
23- service .EXPECT ().BeginTransaction (gomock .Any (), gomock .Any ()).Return (& Ydb_Query.BeginTransactionResponse {
24+ client := NewMockQueryServiceClient (ctrl )
25+ client .EXPECT ().BeginTransaction (gomock .Any (), gomock .Any ()).Return (& Ydb_Query.BeginTransactionResponse {
2426 Status : Ydb .StatusIds_SUCCESS ,
2527 TxMeta : & Ydb_Query.TransactionMeta {
2628 Id : "123" ,
2729 },
2830 }, nil )
2931 t .Log ("begin" )
30- tx , err := begin (ctx , service , & Session {id : "123" }, query .TxSettings ())
32+ tx , err := begin (ctx , client , & Session {id : "123" }, query .TxSettings ())
3133 require .NoError (t , err )
3234 require .Equal (t , "123" , tx .ID ())
3335 })
3436 t .Run ("TransportError" , func (t * testing.T ) {
3537 ctx := xtest .Context (t )
3638 ctrl := gomock .NewController (t )
37- service := NewMockQueryServiceClient (ctrl )
38- service .EXPECT ().BeginTransaction (gomock .Any (), gomock .Any ()).Return (nil , grpcStatus .Error (grpcCodes .Unavailable , "" ))
39+ client := NewMockQueryServiceClient (ctrl )
40+ client .EXPECT ().BeginTransaction (gomock .Any (), gomock .Any ()).Return (nil , grpcStatus .Error (grpcCodes .Unavailable , "" ))
3941 t .Log ("begin" )
40- _ , err := begin (ctx , service , & Session {id : "123" }, query .TxSettings ())
42+ _ , err := begin (ctx , client , & Session {id : "123" }, query .TxSettings ())
4143 require .Error (t , err )
4244 require .True (t , xerrors .IsTransportError (err , grpcCodes .Unavailable ))
4345 })
4446 t .Run ("OperationError" , func (t * testing.T ) {
4547 ctx := xtest .Context (t )
4648 ctrl := gomock .NewController (t )
47- service := NewMockQueryServiceClient (ctrl )
48- service .EXPECT ().BeginTransaction (gomock .Any (), gomock .Any ()).Return (nil ,
49+ client := NewMockQueryServiceClient (ctrl )
50+ client .EXPECT ().BeginTransaction (gomock .Any (), gomock .Any ()).Return (nil ,
4951 xerrors .Operation (xerrors .WithStatusCode (Ydb .StatusIds_UNAVAILABLE )),
5052 )
5153 t .Log ("begin" )
52- _ , err := begin (ctx , service , & Session {id : "123" }, query .TxSettings ())
54+ _ , err := begin (ctx , client , & Session {id : "123" }, query .TxSettings ())
5355 require .Error (t , err )
5456 require .True (t , xerrors .IsOperationError (err , Ydb .StatusIds_UNAVAILABLE ))
5557 })
5658}
59+
60+ func TestCreateSession (t * testing.T ) {
61+ trace := & trace.Query {
62+ OnSessionCreate : func (info trace.QuerySessionCreateStartInfo ) func (info trace.QuerySessionCreateDoneInfo ) {
63+ return func (info trace.QuerySessionCreateDoneInfo ) {
64+ if info .Session != nil && info .Error != nil {
65+ panic ("only one result from tuple may be not nil" )
66+ }
67+ }
68+ },
69+ }
70+ t .Run ("HappyWay" , func (t * testing.T ) {
71+ ctx := xtest .Context (t )
72+ ctrl := gomock .NewController (t )
73+ client := NewMockQueryServiceClient (ctrl )
74+ client .EXPECT ().CreateSession (gomock .Any (), gomock .Any ()).Return (& Ydb_Query.CreateSessionResponse {
75+ Status : Ydb .StatusIds_SUCCESS ,
76+ SessionId : "123" ,
77+ }, nil )
78+ attachStream := NewMockQueryService_AttachSessionClient (ctrl )
79+ attachStream .EXPECT ().Recv ().Return (& Ydb_Query.SessionState {
80+ Status : Ydb .StatusIds_SUCCESS ,
81+ }, nil ).AnyTimes ()
82+ client .EXPECT ().AttachSession (gomock .Any (), & Ydb_Query.AttachSessionRequest {
83+ SessionId : "123" ,
84+ }).Return (attachStream , nil )
85+ t .Log ("createSession" )
86+ require .NotPanics (t , func () {
87+ s , err := createSession (ctx , client , config .New (config .WithTrace (trace )))
88+ require .NoError (t , err )
89+ require .NotNil (t , s )
90+ require .Equal (t , "123" , s .ID ())
91+ })
92+ })
93+ t .Run ("TransportError" , func (t * testing.T ) {
94+ t .Run ("OnCreateSession" , func (t * testing.T ) {
95+ ctx := xtest .Context (t )
96+ ctrl := gomock .NewController (t )
97+ client := NewMockQueryServiceClient (ctrl )
98+ client .EXPECT ().CreateSession (gomock .Any (), gomock .Any ()).Return (nil ,
99+ xerrors .Transport (grpcStatus .Error (grpcCodes .Unavailable , "test" )),
100+ )
101+ t .Log ("createSession" )
102+ require .NotPanics (t , func () {
103+ s , err := createSession (ctx , client , config .New (config .WithTrace (trace )))
104+ require .Error (t , err )
105+ require .Nil (t , s )
106+ })
107+ })
108+ t .Run ("OnAttachStream" , func (t * testing.T ) {
109+ ctx := xtest .Context (t )
110+ ctrl := gomock .NewController (t )
111+ client := NewMockQueryServiceClient (ctrl )
112+ client .EXPECT ().CreateSession (gomock .Any (), gomock .Any ()).Return (& Ydb_Query.CreateSessionResponse {
113+ Status : Ydb .StatusIds_SUCCESS ,
114+ SessionId : "123" ,
115+ }, nil )
116+ client .EXPECT ().AttachSession (gomock .Any (), & Ydb_Query.AttachSessionRequest {
117+ SessionId : "123" ,
118+ }).Return (nil , xerrors .Transport (grpcStatus .Error (grpcCodes .Unavailable , "test" )))
119+ client .EXPECT ().DeleteSession (gomock .Any (), & Ydb_Query.DeleteSessionRequest {
120+ SessionId : "123" ,
121+ }).Return (& Ydb_Query.DeleteSessionResponse {
122+ Status : Ydb .StatusIds_SUCCESS ,
123+ }, nil )
124+ t .Log ("createSession" )
125+ require .NotPanics (t , func () {
126+ s , err := createSession (ctx , client , config .New (config .WithTrace (trace )))
127+ require .Error (t , err )
128+ require .Nil (t , s )
129+ })
130+ })
131+ })
132+ t .Run ("OperationError" , func (t * testing.T ) {
133+ t .Run ("OnCreateSession" , func (t * testing.T ) {
134+ ctx := xtest .Context (t )
135+ ctrl := gomock .NewController (t )
136+ client := NewMockQueryServiceClient (ctrl )
137+ client .EXPECT ().CreateSession (gomock .Any (), gomock .Any ()).Return (nil ,
138+ xerrors .Operation (xerrors .WithStatusCode (Ydb .StatusIds_UNAVAILABLE )),
139+ )
140+ t .Log ("createSession" )
141+ require .NotPanics (t , func () {
142+ s , err := createSession (ctx , client , config .New (config .WithTrace (trace )))
143+ require .Error (t , err )
144+ require .Nil (t , s )
145+ })
146+ })
147+ t .Run ("OnAttachStream" , func (t * testing.T ) {
148+ ctx := xtest .Context (t )
149+ ctrl := gomock .NewController (t )
150+ client := NewMockQueryServiceClient (ctrl )
151+ client .EXPECT ().CreateSession (gomock .Any (), gomock .Any ()).Return (& Ydb_Query.CreateSessionResponse {
152+ Status : Ydb .StatusIds_SUCCESS ,
153+ SessionId : "123" ,
154+ }, nil )
155+ client .EXPECT ().AttachSession (gomock .Any (), & Ydb_Query.AttachSessionRequest {
156+ SessionId : "123" ,
157+ }).Return (nil , xerrors .Operation (xerrors .WithStatusCode (Ydb .StatusIds_UNAVAILABLE )))
158+ client .EXPECT ().DeleteSession (gomock .Any (), & Ydb_Query.DeleteSessionRequest {
159+ SessionId : "123" ,
160+ }).Return (& Ydb_Query.DeleteSessionResponse {
161+ Status : Ydb .StatusIds_SUCCESS ,
162+ }, nil )
163+ t .Log ("createSession" )
164+ require .NotPanics (t , func () {
165+ s , err := createSession (ctx , client , config .New (config .WithTrace (trace )))
166+ require .Error (t , err )
167+ require .Nil (t , s )
168+ })
169+ })
170+ })
171+ }
0 commit comments