@@ -14,6 +14,9 @@ import (
1414 "google.golang.org/grpc/test/bufconn"
1515
1616 "github.com/ydb-platform/ydb-go-sdk/v3/config"
17+ balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
18+ "github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
19+ "github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
1720)
1821
1922func TestBalancer_discoveryConn (t * testing.T ) {
@@ -71,6 +74,52 @@ func TestBalancer_discoveryConn(t *testing.T) {
7174 require .NoError (t , err )
7275}
7376
77+ func TestApplyDiscoveredEndpoints (t * testing.T ) {
78+ ctx := context .Background ()
79+
80+ cfg := config .New ()
81+ pool := conn .NewPool (ctx , cfg )
82+ defer func () { _ = pool .Release (ctx ) }()
83+
84+ b := & Balancer {
85+ driverConfig : cfg ,
86+ pool : pool ,
87+ balancerConfig : balancerConfig.Config {},
88+ }
89+
90+ initial := newConnectionsState (nil , b .balancerConfig .Filter , balancerConfig.Info {}, b .balancerConfig .AllowFallback )
91+ b .connectionsState .Store (initial )
92+
93+ e1 := endpoint .New ("e1.example:2135" , endpoint .WithIPV6 ([]string {"2001:db8::1" }), endpoint .WithID (1 ))
94+ e2 := endpoint .New ("e2.example:2135" , endpoint .WithIPV6 ([]string {"2001:db8::2" }), endpoint .WithID (2 ))
95+
96+ // call with two endpoints
97+ b .applyDiscoveredEndpoints (ctx , []endpoint.Endpoint {e1 , e2 }, "" )
98+
99+ // connectionsState should be updated and reflect the endpoints
100+ after := b .connections ()
101+ require .NotNil (t , after )
102+ all := after .All ()
103+ require .Equal (t , 2 , len (all ))
104+ require .Equal (t , e1 .Address (), all [0 ].Address ())
105+ require .Equal (t , e1 .NodeID (), all [0 ].NodeID ())
106+ require .Equal (t , e2 .Address (), all [1 ].Address ())
107+ require .Equal (t , e2 .NodeID (), all [1 ].NodeID ())
108+
109+ // partially replace endpoints
110+ e3 := endpoint .New ("e3.example:2135" , endpoint .WithIPV6 ([]string {"2001:db8::3" }), endpoint .WithID (1 ))
111+ b .applyDiscoveredEndpoints (ctx , []endpoint.Endpoint {e2 , e3 }, "" )
112+ // connectionsState should be updated and reflect the endpoints
113+ after = b .connections ()
114+ require .NotNil (t , after )
115+ all = after .All ()
116+ require .Equal (t , 2 , len (all ))
117+ require .Equal (t , e2 .Address (), all [0 ].Address ())
118+ require .Equal (t , e2 .NodeID (), all [0 ].NodeID ())
119+ require .Equal (t , e3 .Address (), all [1 ].Address ())
120+ require .Equal (t , e3 .NodeID (), all [1 ].NodeID ())
121+ }
122+
74123// Mock resolver
75124//
76125
0 commit comments