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