tests

Author: asmyasnikov

driver_discovery_client_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
 	"github.com/ydb-platform/ydb-go-sdk/v3/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
 	internalDiscovery "github.com/ydb-platform/ydb-go-sdk/v3/internal/discovery"
@@ -40,6 +41,40 @@ func TestDriverDiscoveryClientCloseReleasesBootstrapRef(t *testing.T) {
 	require.NotSame(t, bootstrap, again)
 }
 
+func TestDriverConnectInitializesDiscoveryClient(t *testing.T) {
+	ctx := context.Background()
+	cfg := config.New(
+		config.WithEndpoint("bootstrap:2135"),
+		config.WithDatabase("/local"),
+		config.WithBalancer(balancers.SingleConn()),
+	)
+	pool := conn.NewPool(ctx, cfg)
+
+	d := &Driver{
+		config: cfg,
+		pool:   pool,
+		metaBalancer: &balancerWithMeta{
+			meta: cfg.Meta(),
+			close: func(context.Context) error {
+				return nil
+			},
+		},
+	}
+
+	require.NoError(t, d.connect(ctx))
+	t.Cleanup(func() {
+		require.NoError(t, d.metaBalancer.Close(ctx))
+	})
+
+	discoveryClient, err := d.discovery.Get()
+	require.NoError(t, err)
+	require.NotNil(t, discoveryClient)
+	require.NotNil(t, discoveryClient.conn)
+
+	require.NoError(t, discoveryClient.Close(ctx))
+	require.NoError(t, pool.RemoveRef(ctx))
+}
+
 func TestDriverDiscoveryInitFailsWhenPoolClosed(t *testing.T) {
 	ctx := context.Background()
 	pool := conn.NewPool(ctx, config.New())

internal/conn/pool.go

@@ -267,6 +267,18 @@ func (p *Pool) closeConnsForFailedResolve(ctx context.Context, target string) {
 	}
 }
 
+func (p *Pool) onResolveCallback(
+	ctx context.Context,
+	start trace.DriverResolveStartInfo,
+) func(trace.DriverResolveDoneInfo) {
+	target := start.Target
+	resolved := start.Resolved
+
+	return func(info trace.DriverResolveDoneInfo) {
+		p.onResolveDone(ctx, target, resolved, info)
+	}
+}
+
 func (p *Pool) onResolveDone(ctx context.Context, target string, resolved []string, info trace.DriverResolveDoneInfo) {
 	if info.Error != nil || len(resolved) == 0 {
 		// Reset gRPC transport only; keep map entries and useCount unchanged.
@@ -294,12 +306,7 @@ func NewPool(ctx context.Context, config Config) *Pool {
 		grpc.WithResolvers(
 			xresolver.New("", gtrace.Compose(config.Trace(), &trace.Driver{
 				OnResolve: func(info trace.DriverResolveStartInfo) func(trace.DriverResolveDoneInfo) {
-					target := info.Target
-					resolved := info.Resolved
-
-					return func(info trace.DriverResolveDoneInfo) {
-						p.onResolveDone(ctx, target, resolved, info)
-					}
+					return p.onResolveCallback(ctx, info)
 				},
 			})),
 		),

internal/conn/pool_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn/state"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
@@ -999,3 +1000,99 @@ func TestPool_EndpointsToConnectionsNilMap(t *testing.T) {
 	e := endpoint.New("closed:2135")
 	require.Empty(t, endpointsToConnections(pool, []endpoint.Endpoint{e}))
 }
+
+func TestPool_AddRefReturnsErrorWhenClosed(t *testing.T) {
+	ctx := context.Background()
+	pool := NewPool(ctx, &mockConfig{})
+	require.NoError(t, pool.RemoveRef(ctx))
+
+	err := pool.AddRef(ctx)
+	require.ErrorIs(t, err, ErrClosedPool)
+}
+
+func TestPool_RemoveRefReturnsCloseErrors(t *testing.T) {
+	ctx := context.Background()
+	pool := NewPool(ctx, &mockConfig{})
+
+	e := endpoint.New("close-error:2135")
+	cc := newConn(e, pool)
+	cc.mtx.Lock()
+	cc.grpcConn = &mockGrpcConn{closeErr: errors.New("grpc close failed")}
+	cc.mtx.Unlock()
+
+	pool.mu.Lock()
+	pool.conns[e.Key()] = &connValue{cc: cc}
+	pool.mu.Unlock()
+
+	err := pool.RemoveRef(ctx)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "connection pool close failed")
+}
+
+func TestPool_OnResolveCallback(t *testing.T) {
+	ctx := context.Background()
+	pool := NewPool(ctx, &mockConfig{})
+	defer func() {
+		_ = pool.RemoveRef(ctx)
+	}()
+
+	target := "localhost:2135"
+	e := endpoint.New(target)
+	_ = pool.Get(e)
+
+	done := pool.onResolveCallback(ctx, trace.DriverResolveStartInfo{
+		Target:   target,
+		Resolved: []string{target},
+	})
+	done(trace.DriverResolveDoneInfo{Error: errors.New("resolve failed")})
+
+	require.True(t, testPoolHasConn(pool, e.Key()))
+}
+
+func TestPool_CloseConnsForFailedResolveSkipsAlreadyClosed(t *testing.T) {
+	ctx := context.Background()
+	pool := NewPool(ctx, &mockConfig{})
+	defer func() {
+		_ = pool.RemoveRef(ctx)
+	}()
+
+	target := "localhost:2135"
+	e := endpoint.New(target)
+	cc := newConn(e, pool)
+	cc.mtx.Lock()
+	cc.closed = true
+	cc.mtx.Unlock()
+
+	pool.mu.Lock()
+	pool.conns[e.Key()] = &connValue{cc: cc}
+	pool.mu.Unlock()
+
+	require.NotPanics(t, func() {
+		pool.closeConnsForFailedResolve(ctx, target)
+	})
+}
+
+type mockGrpcConn struct {
+	closeErr error
+}
+
+func (m *mockGrpcConn) Invoke(context.Context, string, any, any, ...grpc.CallOption) error {
+	return nil
+}
+
+func (m *mockGrpcConn) NewStream(
+	context.Context,
+	*grpc.StreamDesc,
+	string,
+	...grpc.CallOption,
+) (grpc.ClientStream, error) {
+	return nil, errors.New("not implemented")
+}
+
+func (m *mockGrpcConn) Close() error {
+	return m.closeErr
+}
+
+func (m *mockGrpcConn) GetState() connectivity.State {
+	return connectivity.Ready
+}