refactor(balancer): update ClientConn to use ClientConnWrapper

pkg/balancer/balancer.go

@@ -18,7 +18,7 @@ type Balancer struct {
 	state connectivity.State
 
 	// The ClientConn to communicate with the gRPC client.
-	clientConn balancer.ClientConn
+	clientConn ClientConnWrapper
 
 	// Current picker used to select SubConns for requests.
 	picker balancer.Picker

pkg/balancer/balancer_test.go

@@ -4,6 +4,8 @@ import (
 	"errors"
 	"testing"
 
+	estats "google.golang.org/grpc/experimental/stats"
+
 	"github.com/cespare/xxhash/v2"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
@@ -27,6 +29,10 @@ type mockClientConn struct {
 	newSubConn  balancer.SubConn
 }
 
+func (m *mockClientConn) MetricsRecorder() estats.MetricsRecorder {
+	return nil
+}
+
 func (m *mockClientConn) UpdateState(state balancer.State) {
 	m.updateState = state
 }

pkg/balancer/wrapper.go

@@ -0,0 +1,25 @@
+package balancer
+
+import (
+	"google.golang.org/grpc/balancer"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/resolver"
+)
+
+// ClientConnWrapper is an interface that wraps the gRPC ClientConn interface.
+type ClientConnWrapper interface {
+	// NewSubConn creates a new SubConn with the specified addresses and options.
+	NewSubConn([]resolver.Address, balancer.NewSubConnOptions) (balancer.SubConn, error)
+	// RemoveSubConn removes the specified SubConn.
+	RemoveSubConn(balancer.SubConn)
+	// UpdateAddresses updates the addresses for the specified SubConn.
+	UpdateAddresses(balancer.SubConn, []resolver.Address)
+	// UpdateState updates the overall connectivity state of the balancer.
+	UpdateState(balancer.State)
+	// ResolveNow triggers an immediate resolution of the target.
+	ResolveNow(resolver.ResolveNowOptions)
+	// Target returns the target URI of the connection.
+	Target() string
+	// MetricsRecorder returns the metrics recorder for the connection.
+	MetricsRecorder() estats.MetricsRecorder
+}