github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/balancers/balancers_test.go (about)

     1  package balancers
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/require"
     7  
     8  	balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
    10  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/mock"
    12  )
    13  
    14  func TestPreferLocalDC(t *testing.T) {
    15  	conns := []conn.Conn{
    16  		&mock.Conn{AddrField: "1", LocationField: "1"},
    17  		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "2"},
    18  		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "2"},
    19  	}
    20  	rr := PreferNearestDC(RandomChoice())
    21  	require.False(t, rr.AllowFallback)
    22  	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(balancerConfig.Info{SelfLocation: "2"}, rr, conns))
    23  }
    24  
    25  func TestPreferLocalDCWithFallBack(t *testing.T) {
    26  	conns := []conn.Conn{
    27  		&mock.Conn{AddrField: "1", LocationField: "1"},
    28  		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "2"},
    29  		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "2"},
    30  	}
    31  	rr := PreferNearestDCWithFallBack(RandomChoice())
    32  	require.True(t, rr.AllowFallback)
    33  	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(balancerConfig.Info{SelfLocation: "2"}, rr, conns))
    34  }
    35  
    36  func TestPreferLocations(t *testing.T) {
    37  	conns := []conn.Conn{
    38  		&mock.Conn{AddrField: "1", LocationField: "zero", State: conn.Online},
    39  		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "one"},
    40  		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "two"},
    41  	}
    42  
    43  	rr := PreferLocations(RandomChoice(), "zero", "two")
    44  	require.False(t, rr.AllowFallback)
    45  	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(balancerConfig.Info{}, rr, conns))
    46  }
    47  
    48  func TestPreferLocationsWithFallback(t *testing.T) {
    49  	conns := []conn.Conn{
    50  		&mock.Conn{AddrField: "1", LocationField: "zero", State: conn.Online},
    51  		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "one"},
    52  		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "two"},
    53  	}
    54  
    55  	rr := PreferLocationsWithFallback(RandomChoice(), "zero", "two")
    56  	require.True(t, rr.AllowFallback)
    57  	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(balancerConfig.Info{}, rr, conns))
    58  }
    59  
    60  func applyPreferFilter(info balancerConfig.Info, b *balancerConfig.Config, conns []conn.Conn) []conn.Conn {
    61  	if b.Filter == nil {
    62  		b.Filter = filterFunc(func(info balancerConfig.Info, e endpoint.Info) bool { return true })
    63  	}
    64  	res := make([]conn.Conn, 0, len(conns))
    65  	for _, c := range conns {
    66  		if b.Filter.Allow(info, c.Endpoint()) {
    67  			res = append(res, c)
    68  		}
    69  	}
    70  
    71  	return res
    72  }