github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/tests/integration/discovery_test.go (about)

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  	"crypto/tls"
     9  	"os"
    10  	"testing"
    11  	"time"
    12  
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/metadata"
    15  
    16  	"github.com/ydb-platform/ydb-go-sdk/v3"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/config"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
    20  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    21  	"github.com/ydb-platform/ydb-go-sdk/v3/log"
    22  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    23  )
    24  
    25  func TestDiscovery(t *testing.T) {
    26  	var (
    27  		userAgent     = "connection user agent"
    28  		requestType   = "connection request type"
    29  		checkMedatada = func(ctx context.Context) {
    30  			md, has := metadata.FromOutgoingContext(ctx)
    31  			if !has {
    32  				t.Fatalf("no medatada")
    33  			}
    34  			userAgents := md.Get(meta.HeaderUserAgent)
    35  			if len(userAgents) == 0 {
    36  				t.Fatalf("no user agent")
    37  			}
    38  			if userAgents[0] != userAgent {
    39  				t.Fatalf("unknown user agent: %s", userAgents[0])
    40  			}
    41  			requestTypes := md.Get(meta.HeaderRequestType)
    42  			if len(requestTypes) == 0 {
    43  				t.Fatalf("no request type")
    44  			}
    45  			if requestTypes[0] != requestType {
    46  				t.Fatalf("unknown request type: %s", requestTypes[0])
    47  			}
    48  		}
    49  		parking = make(chan struct{})
    50  		ctx     = xtest.Context(t)
    51  	)
    52  
    53  	db, err := ydb.Open(ctx,
    54  		os.Getenv("YDB_CONNECTION_STRING"),
    55  		ydb.WithAccessTokenCredentials(
    56  			os.Getenv("YDB_ACCESS_TOKEN_CREDENTIALS"),
    57  		),
    58  		ydb.With(
    59  			config.WithOperationTimeout(time.Second*2),
    60  			config.WithOperationCancelAfter(time.Second*2),
    61  		),
    62  		ydb.WithBalancer(balancers.SingleConn()),
    63  		ydb.WithConnectionTTL(time.Second*1),
    64  		ydb.WithMinTLSVersion(tls.VersionTLS10),
    65  		ydb.WithLogger(
    66  			newLoggerWithMinLevel(t, log.WARN),
    67  			trace.MatchDetails(`ydb\.(driver|discovery|repeater).*`),
    68  		),
    69  		ydb.WithUserAgent(userAgent),
    70  		ydb.WithRequestsType(requestType),
    71  		ydb.With(
    72  			config.WithGrpcOptions(
    73  				grpc.WithUnaryInterceptor(func(
    74  					ctx context.Context,
    75  					method string,
    76  					req, reply interface{},
    77  					cc *grpc.ClientConn,
    78  					invoker grpc.UnaryInvoker,
    79  					opts ...grpc.CallOption,
    80  				) error {
    81  					checkMedatada(ctx)
    82  					return invoker(ctx, method, req, reply, cc, opts...)
    83  				}),
    84  				grpc.WithStreamInterceptor(func(
    85  					ctx context.Context,
    86  					desc *grpc.StreamDesc,
    87  					cc *grpc.ClientConn,
    88  					method string,
    89  					streamer grpc.Streamer,
    90  					opts ...grpc.CallOption,
    91  				) (grpc.ClientStream, error) {
    92  					checkMedatada(ctx)
    93  					return streamer(ctx, desc, cc, method, opts...)
    94  				}),
    95  			),
    96  		),
    97  		ydb.WithTraceDriver(trace.Driver{
    98  			OnConnPark: func(info trace.DriverConnParkStartInfo) func(trace.DriverConnParkDoneInfo) {
    99  				return func(info trace.DriverConnParkDoneInfo) {
   100  					parking <- struct{}{}
   101  				}
   102  			},
   103  		}),
   104  	)
   105  	if err != nil {
   106  		t.Fatal(err)
   107  	}
   108  	defer func() {
   109  		// cleanup connection
   110  		if e := db.Close(ctx); e != nil {
   111  			t.Fatalf("db close failed: %+v", e)
   112  		}
   113  	}()
   114  	t.Run("discovery.Discover", func(t *testing.T) {
   115  		if endpoints, err := db.Discovery().Discover(ctx); err != nil {
   116  			t.Fatal(err)
   117  		} else {
   118  			t.Log(endpoints)
   119  		}
   120  		t.Run("wait", func(t *testing.T) {
   121  			t.Run("parking", func(t *testing.T) {
   122  				<-parking // wait for parking conn
   123  				t.Run("re-discover", func(t *testing.T) {
   124  					if endpoints, err := db.Discovery().Discover(ctx); err != nil {
   125  						t.Fatal(err)
   126  					} else {
   127  						t.Log(endpoints)
   128  					}
   129  				})
   130  			})
   131  		})
   132  	})
   133  }