From bca56ef8c352ac81ae1d98f52d879cf99a8c863e Mon Sep 17 00:00:00 2001
From: kevin <wanjunfeng@gmail.com>
Date: Fri, 7 Aug 2020 15:28:40 +0800
Subject: [PATCH] add tests

---
 core/mathx/entropy.go                  | 10 ++-
 rpcx/internal/balancer/p2c/p2c.go      |  1 -
 rpcx/internal/balancer/p2c/p2c_test.go | 87 +++++++++++++++++++++++++-
 3 files changed, 95 insertions(+), 3 deletions(-)

diff --git a/core/mathx/entropy.go b/core/mathx/entropy.go
index f426f2c8..e1b589b8 100644
--- a/core/mathx/entropy.go
+++ b/core/mathx/entropy.go
@@ -2,9 +2,14 @@ package mathx
 
 import "math"
 
+const epsilon = 1e-6
+
 func CalcEntropy(m map[interface{}]int) float64 {
-	var entropy float64
+	if len(m) == 0 || len(m) == 1 {
+		return 1
+	}
 
+	var entropy float64
 	var total int
 	for _, v := range m {
 		total += v
@@ -12,6 +17,9 @@ func CalcEntropy(m map[interface{}]int) float64 {
 
 	for _, v := range m {
 		proba := float64(v) / float64(total)
+		if proba < epsilon {
+			proba = epsilon
+		}
 		entropy -= proba * math.Log2(proba)
 	}
 
diff --git a/rpcx/internal/balancer/p2c/p2c.go b/rpcx/internal/balancer/p2c/p2c.go
index 8194b858..1bea4742 100644
--- a/rpcx/internal/balancer/p2c/p2c.go
+++ b/rpcx/internal/balancer/p2c/p2c.go
@@ -169,7 +169,6 @@ func (p *p2cPicker) logStats() {
 	defer p.lock.Unlock()
 
 	for _, conn := range p.conns {
-		fmt.Println(conn.lag, conn.inflight)
 		stats = append(stats, fmt.Sprintf("conn: %s, load: %d, reqs: %d",
 			conn.addr.Addr, conn.load(), atomic.SwapInt64(&conn.requests, 0)))
 	}
diff --git a/rpcx/internal/balancer/p2c/p2c_test.go b/rpcx/internal/balancer/p2c/p2c_test.go
index 685780df..b3a1528d 100644
--- a/rpcx/internal/balancer/p2c/p2c_test.go
+++ b/rpcx/internal/balancer/p2c/p2c_test.go
@@ -1,7 +1,92 @@
 package p2c
 
-import "testing"
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"testing"
+
+	"zero/core/logx"
+	"zero/core/mathx"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/status"
+)
+
+func init() {
+	logx.Disable()
+}
+
+func TestP2cPicker_PickNil(t *testing.T) {
+	builder := new(p2cPickerBuilder)
+	picker := builder.Build(nil)
+	_, _, err := picker.Pick(context.Background(), balancer.PickInfo{
+		FullMethodName: "/",
+		Ctx:            context.Background(),
+	})
+	assert.NotNil(t, err)
+}
 
 func TestP2cPicker_Pick(t *testing.T) {
+	tests := []struct {
+		name       string
+		candidates int
+	}{
+		{
+			name:       "single",
+			candidates: 1,
+		},
+		{
+			name:       "multiple",
+			candidates: 100,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			builder := new(p2cPickerBuilder)
+			ready := make(map[resolver.Address]balancer.SubConn)
+			for i := 0; i < test.candidates; i++ {
+				ready[resolver.Address{
+					Addr: strconv.Itoa(i),
+				}] = new(mockClientConn)
+			}
+
+			picker := builder.Build(ready)
+			for i := 0; i < 10000; i++ {
+				_, done, err := picker.Pick(context.Background(), balancer.PickInfo{
+					FullMethodName: "/",
+					Ctx:            context.Background(),
+				})
+				assert.Nil(t, err)
+				if i%100 == 0 {
+					err = status.Error(codes.DeadlineExceeded, "deadline")
+				}
+				done(balancer.DoneInfo{
+					Err: err,
+				})
+			}
+
+			dist := make(map[interface{}]int)
+			conns := picker.(*p2cPicker).conns
+			for _, conn := range conns {
+				dist[conn.addr.Addr] = int(conn.requests)
+			}
+
+			entropy := mathx.CalcEntropy(dist)
+			assert.True(t, entropy > .95, fmt.Sprintf("entropy is %f, less than .95", entropy))
+		})
+	}
+}
+
+type mockClientConn struct {
+}
+
+func (m mockClientConn) UpdateAddresses(addresses []resolver.Address) {
+}
 
+func (m mockClientConn) Connect() {
 }