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

     1  package discovery
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/jonboulle/clockwork"
     7  	"github.com/stretchr/testify/require"
     8  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
     9  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Discovery"
    10  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Operations"
    11  	"go.uber.org/mock/gomock"
    12  	grpcCodes "google.golang.org/grpc/codes"
    13  	"google.golang.org/grpc/status"
    14  	"google.golang.org/protobuf/types/known/anypb"
    15  
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/discovery/config"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    20  )
    21  
    22  func TestDiscover(t *testing.T) {
    23  	t.Run("HappyWay", func(t *testing.T) {
    24  		ctx := xtest.Context(t)
    25  		ctrl := gomock.NewController(t)
    26  		clock := clockwork.NewFakeClock()
    27  		client := NewMockDiscoveryServiceClient(ctrl)
    28  		client.EXPECT().ListEndpoints(gomock.Any(), &Ydb_Discovery.ListEndpointsRequest{
    29  			Database: "test",
    30  		}).Return(&Ydb_Discovery.ListEndpointsResponse{
    31  			Operation: &Ydb_Operations.Operation{
    32  				Ready:  true,
    33  				Status: Ydb.StatusIds_SUCCESS,
    34  				Result: xtest.Must(anypb.New(&Ydb_Discovery.ListEndpointsResult{
    35  					Endpoints: []*Ydb_Discovery.EndpointInfo{
    36  						{
    37  							Address: "node1",
    38  							Port:    1,
    39  							Ssl:     true,
    40  						},
    41  						{
    42  							Address:  "node2",
    43  							Port:     2,
    44  							Location: "AZ0",
    45  							Ssl:      true,
    46  						},
    47  						{
    48  							Address: "node3",
    49  							Port:    3,
    50  							Ssl:     false,
    51  						},
    52  						{
    53  							Address:  "node4",
    54  							Port:     4,
    55  							Location: "AZ0",
    56  							Ssl:      false,
    57  						},
    58  					},
    59  					SelfLocation: "AZ0",
    60  				})),
    61  			},
    62  		}, nil)
    63  		endpoints, location, err := discover(ctx, client, config.New(
    64  			config.WithDatabase("test"),
    65  			config.WithSecure(false),
    66  			config.WithClock(clock),
    67  		))
    68  		require.NoError(t, err)
    69  		require.EqualValues(t, "AZ0", location)
    70  		require.EqualValues(t, []endpoint.Endpoint{
    71  			endpoint.New("node3:3",
    72  				endpoint.WithLocalDC(false),
    73  				endpoint.WithLastUpdated(clock.Now()),
    74  			),
    75  			endpoint.New("node4:4",
    76  				endpoint.WithLocalDC(true),
    77  				endpoint.WithLocation("AZ0"),
    78  				endpoint.WithLastUpdated(clock.Now()),
    79  			),
    80  		}, endpoints)
    81  	})
    82  	t.Run("TransportError", func(t *testing.T) {
    83  		ctx := xtest.Context(t)
    84  		ctrl := gomock.NewController(t)
    85  		client := NewMockDiscoveryServiceClient(ctrl)
    86  		client.EXPECT().ListEndpoints(gomock.Any(), &Ydb_Discovery.ListEndpointsRequest{
    87  			Database: "test",
    88  		}).Return(nil, xerrors.Transport(status.Error(grpcCodes.Unavailable, "")))
    89  		endpoints, location, err := discover(ctx, client, config.New(
    90  			config.WithDatabase("test"),
    91  		))
    92  		require.Error(t, err)
    93  		require.Empty(t, endpoints)
    94  		require.Equal(t, "", location)
    95  		require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
    96  	})
    97  	t.Run("OperationError", func(t *testing.T) {
    98  		ctx := xtest.Context(t)
    99  		ctrl := gomock.NewController(t)
   100  		client := NewMockDiscoveryServiceClient(ctrl)
   101  		client.EXPECT().ListEndpoints(gomock.Any(), &Ydb_Discovery.ListEndpointsRequest{
   102  			Database: "test",
   103  		}).Return(&Ydb_Discovery.ListEndpointsResponse{
   104  			Operation: &Ydb_Operations.Operation{
   105  				Ready:  true,
   106  				Status: Ydb.StatusIds_UNAVAILABLE,
   107  			},
   108  		}, nil)
   109  		endpoints, location, err := discover(ctx, client, config.New(
   110  			config.WithDatabase("test"),
   111  		))
   112  		require.Error(t, err)
   113  		require.Empty(t, endpoints)
   114  		require.Equal(t, "", location)
   115  		require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
   116  	})
   117  	t.Run("WithAddressMutator", func(t *testing.T) {
   118  		ctx := xtest.Context(t)
   119  		ctrl := gomock.NewController(t)
   120  		clock := clockwork.NewFakeClock()
   121  		client := NewMockDiscoveryServiceClient(ctrl)
   122  		client.EXPECT().ListEndpoints(gomock.Any(), &Ydb_Discovery.ListEndpointsRequest{
   123  			Database: "test",
   124  		}).Return(&Ydb_Discovery.ListEndpointsResponse{
   125  			Operation: &Ydb_Operations.Operation{
   126  				Ready:  true,
   127  				Status: Ydb.StatusIds_SUCCESS,
   128  				Result: xtest.Must(anypb.New(&Ydb_Discovery.ListEndpointsResult{
   129  					Endpoints: []*Ydb_Discovery.EndpointInfo{
   130  						{
   131  							Address: "node1",
   132  							Port:    1,
   133  						},
   134  						{
   135  							Address:  "node2",
   136  							Port:     2,
   137  							Location: "AZ0",
   138  						},
   139  					},
   140  					SelfLocation: "AZ0",
   141  				})),
   142  			},
   143  		}, nil)
   144  		endpoints, location, err := discover(ctx, client, config.New(
   145  			config.WithDatabase("test"),
   146  			config.WithAddressMutator(func(address string) string {
   147  				return "u-" + address
   148  			}),
   149  			config.WithClock(clock),
   150  		))
   151  		require.NoError(t, err)
   152  		require.EqualValues(t, "AZ0", location)
   153  		require.EqualValues(t, []endpoint.Endpoint{
   154  			endpoint.New("u-node1:1",
   155  				endpoint.WithLocalDC(false),
   156  				endpoint.WithLastUpdated(clock.Now()),
   157  			),
   158  			endpoint.New("u-node2:2",
   159  				endpoint.WithLocalDC(true),
   160  				endpoint.WithLocation("AZ0"),
   161  				endpoint.WithLastUpdated(clock.Now()),
   162  			),
   163  		}, endpoints)
   164  	})
   165  }