test case for balancer

app/router/config.go

@@ -132,5 +132,6 @@ func (br *BalancingRule) Build(ohm outbound.Manager) (*Balancer, error) {
 	return &Balancer{
 		selectors: br.OutboundSelector,
 		strategy:  &RandomStrategy{},
+		ohm:       ohm,
 	}, nil
 }

app/router/router_test.go

@@ -8,9 +8,15 @@ import (
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/session"
+	"v2ray.com/core/features/outbound"
 	"v2ray.com/core/testing/mocks"
 )
 
+type mockOutboundManager struct {
+	outbound.Manager
+	outbound.HandlerSelector
+}
+
 func TestSimpleRouter(t *testing.T) {
 	config := &Config{
 		Rule: []*RoutingRule{
@@ -29,9 +35,57 @@ func TestSimpleRouter(t *testing.T) {
 	defer mockCtl.Finish()
 
 	mockDns := mocks.NewDNSClient(mockCtl)
+	mockOhm := mocks.NewOutboundManager(mockCtl)
+	mockHs := mocks.NewOutboundHandlerSelector(mockCtl)
 
 	r := new(Router)
-	common.Must(r.Init(config, mockDns, nil))
+	common.Must(r.Init(config, mockDns, &mockOutboundManager{
+		Manager:         mockOhm,
+		HandlerSelector: mockHs,
+	}))
+
+	ctx := withOutbound(&session.Outbound{Target: net.TCPDestination(net.DomainAddress("v2ray.com"), 80)})
+	tag, err := r.PickRoute(ctx)
+	common.Must(err)
+	if tag != "test" {
+		t.Error("expect tag 'test', bug actually ", tag)
+	}
+}
+
+func TestSimpleBalancer(t *testing.T) {
+	config := &Config{
+		Rule: []*RoutingRule{
+			{
+				TargetTag: &RoutingRule_BalancingTag{
+					BalancingTag: "balance",
+				},
+				NetworkList: &net.NetworkList{
+					Network: []net.Network{net.Network_TCP},
+				},
+			},
+		},
+		BalancingRule: []*BalancingRule{
+			{
+				Tag:              "balance",
+				OutboundSelector: []string{"test-"},
+			},
+		},
+	}
+
+	mockCtl := gomock.NewController(t)
+	defer mockCtl.Finish()
+
+	mockDns := mocks.NewDNSClient(mockCtl)
+	mockOhm := mocks.NewOutboundManager(mockCtl)
+	mockHs := mocks.NewOutboundHandlerSelector(mockCtl)
+
+	mockHs.EXPECT().Select(gomock.Eq([]string{"test-"})).Return([]string{"test"})
+
+	r := new(Router)
+	common.Must(r.Init(config, mockDns, &mockOutboundManager{
+		Manager:         mockOhm,
+		HandlerSelector: mockHs,
+	}))
 
 	ctx := withOutbound(&session.Outbound{Target: net.TCPDestination(net.DomainAddress("v2ray.com"), 80)})
 	tag, err := r.PickRoute(ctx)

mocks.go

@@ -6,4 +6,5 @@ package core
 //go:generate mockgen -package mocks -destination testing/mocks/io.go -mock_names Reader=Reader,Writer=Writer io Reader,Writer
 //go:generate mockgen -package mocks -destination testing/mocks/mux.go -mock_names ClientWorkerFactory=MuxClientWorkerFactory v2ray.com/core/common/mux ClientWorkerFactory
 //go:generate mockgen -package mocks -destination testing/mocks/dns.go -mock_names Client=DNSClient v2ray.com/core/features/dns Client
+//go:generate mockgen -package mocks -destination testing/mocks/outbound.go -mock_names Manager=OutboundManager,HandlerSelector=OutboundHandlerSelector v2ray.com/core/features/outbound Manager,HandlerSelector
 //go:generate mockgen -package mocks -destination testing/mocks/proxy.go -mock_names Inbound=ProxyInbound,Outbound=ProxyOutbound v2ray.com/core/proxy Inbound,Outbound

testing/mocks/outbound.go

@@ -0,0 +1,154 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: v2ray.com/core/features/outbound (interfaces: Manager,HandlerSelector)
+
+// Package mocks is a generated GoMock package.
+package mocks
+
+import (
+	context "context"
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+	outbound "v2ray.com/core/features/outbound"
+)
+
+// OutboundManager is a mock of Manager interface
+type OutboundManager struct {
+	ctrl     *gomock.Controller
+	recorder *OutboundManagerMockRecorder
+}
+
+// OutboundManagerMockRecorder is the mock recorder for OutboundManager
+type OutboundManagerMockRecorder struct {
+	mock *OutboundManager
+}
+
+// NewOutboundManager creates a new mock instance
+func NewOutboundManager(ctrl *gomock.Controller) *OutboundManager {
+	mock := &OutboundManager{ctrl: ctrl}
+	mock.recorder = &OutboundManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *OutboundManager) EXPECT() *OutboundManagerMockRecorder {
+	return m.recorder
+}
+
+// AddHandler mocks base method
+func (m *OutboundManager) AddHandler(arg0 context.Context, arg1 outbound.Handler) error {
+	ret := m.ctrl.Call(m, "AddHandler", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// AddHandler indicates an expected call of AddHandler
+func (mr *OutboundManagerMockRecorder) AddHandler(arg0, arg1 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddHandler", reflect.TypeOf((*OutboundManager)(nil).AddHandler), arg0, arg1)
+}
+
+// Close mocks base method
+func (m *OutboundManager) Close() error {
+	ret := m.ctrl.Call(m, "Close")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Close indicates an expected call of Close
+func (mr *OutboundManagerMockRecorder) Close() *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*OutboundManager)(nil).Close))
+}
+
+// GetDefaultHandler mocks base method
+func (m *OutboundManager) GetDefaultHandler() outbound.Handler {
+	ret := m.ctrl.Call(m, "GetDefaultHandler")
+	ret0, _ := ret[0].(outbound.Handler)
+	return ret0
+}
+
+// GetDefaultHandler indicates an expected call of GetDefaultHandler
+func (mr *OutboundManagerMockRecorder) GetDefaultHandler() *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDefaultHandler", reflect.TypeOf((*OutboundManager)(nil).GetDefaultHandler))
+}
+
+// GetHandler mocks base method
+func (m *OutboundManager) GetHandler(arg0 string) outbound.Handler {
+	ret := m.ctrl.Call(m, "GetHandler", arg0)
+	ret0, _ := ret[0].(outbound.Handler)
+	return ret0
+}
+
+// GetHandler indicates an expected call of GetHandler
+func (mr *OutboundManagerMockRecorder) GetHandler(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHandler", reflect.TypeOf((*OutboundManager)(nil).GetHandler), arg0)
+}
+
+// RemoveHandler mocks base method
+func (m *OutboundManager) RemoveHandler(arg0 context.Context, arg1 string) error {
+	ret := m.ctrl.Call(m, "RemoveHandler", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// RemoveHandler indicates an expected call of RemoveHandler
+func (mr *OutboundManagerMockRecorder) RemoveHandler(arg0, arg1 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveHandler", reflect.TypeOf((*OutboundManager)(nil).RemoveHandler), arg0, arg1)
+}
+
+// Start mocks base method
+func (m *OutboundManager) Start() error {
+	ret := m.ctrl.Call(m, "Start")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Start indicates an expected call of Start
+func (mr *OutboundManagerMockRecorder) Start() *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Start", reflect.TypeOf((*OutboundManager)(nil).Start))
+}
+
+// Type mocks base method
+func (m *OutboundManager) Type() interface{} {
+	ret := m.ctrl.Call(m, "Type")
+	ret0, _ := ret[0].(interface{})
+	return ret0
+}
+
+// Type indicates an expected call of Type
+func (mr *OutboundManagerMockRecorder) Type() *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Type", reflect.TypeOf((*OutboundManager)(nil).Type))
+}
+
+// OutboundHandlerSelector is a mock of HandlerSelector interface
+type OutboundHandlerSelector struct {
+	ctrl     *gomock.Controller
+	recorder *OutboundHandlerSelectorMockRecorder
+}
+
+// OutboundHandlerSelectorMockRecorder is the mock recorder for OutboundHandlerSelector
+type OutboundHandlerSelectorMockRecorder struct {
+	mock *OutboundHandlerSelector
+}
+
+// NewOutboundHandlerSelector creates a new mock instance
+func NewOutboundHandlerSelector(ctrl *gomock.Controller) *OutboundHandlerSelector {
+	mock := &OutboundHandlerSelector{ctrl: ctrl}
+	mock.recorder = &OutboundHandlerSelectorMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *OutboundHandlerSelector) EXPECT() *OutboundHandlerSelectorMockRecorder {
+	return m.recorder
+}
+
+// Select mocks base method
+func (m *OutboundHandlerSelector) Select(arg0 []string) []string {
+	ret := m.ctrl.Call(m, "Select", arg0)
+	ret0, _ := ret[0].([]string)
+	return ret0
+}
+
+// Select indicates an expected call of Select
+func (mr *OutboundHandlerSelectorMockRecorder) Select(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Select", reflect.TypeOf((*OutboundHandlerSelector)(nil).Select), arg0)
+}