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)