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