Add multiple connection tests

bojand · Apr 14, 2019 · 1d85a9f · 1d85a9f
1 parent 832ca2c
commit 1d85a9f
Show file tree

Hide file tree

Showing 4 changed files with 203 additions and 1 deletion.
diff --git a/internal/common.go b/internal/common.go
@@ -36,12 +36,18 @@ func StartServer(secure bool) (*helloworld.Greeter, *grpc.Server, error) {
 		opts = []grpc.ServerOption{grpc.Creds(creds)}
 	}
 
+	stats := helloworld.NewHWStats()
+
+	opts = append(opts, grpc.StatsHandler(stats))
+
 	s := grpc.NewServer(opts...)
 
 	gs := helloworld.NewGreeter()
 	helloworld.RegisterGreeterServer(s, gs)
 	reflection.Register(s)
 
+	gs.Stats = stats
+
 	TestPort = strconv.Itoa(lis.Addr().(*net.TCPAddr).Port)
 	TestLocalhost = "localhost:" + TestPort
 

diff --git a/internal/helloworld/greeter_server.go b/internal/helloworld/greeter_server.go
@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	context "golang.org/x/net/context"
+	"google.golang.org/grpc/stats"
 )
 
 // CallType represents one of the gRPC call types:
@@ -27,8 +28,10 @@ var Bidi CallType = "bidi"
 // Greeter implements the GreeterServer for tests
 type Greeter struct {
 	streamData []*HelloReply
-	mutex      *sync.RWMutex
 
+	Stats *HWStatsHandler
+
+	mutex      *sync.RWMutex
 	callCounts map[CallType]int
 }
 
@@ -107,6 +110,12 @@ func (s *Greeter) ResetCounters() {
 	s.callCounts[ClientStream] = 0
 	s.callCounts[Bidi] = 0
 	s.mutex.Unlock()
+
+	if s.Stats != nil {
+		s.Stats.mutex.Lock()
+		s.Stats.connCount = 0
+		s.Stats.mutex.Unlock()
+	}
 }
 
 // GetCount gets the count for specific call type
@@ -120,6 +129,11 @@ func (s *Greeter) GetCount(key CallType) int {
 	return -1
 }
 
+// GetConnectionCount gets the connection count
+func (s *Greeter) GetConnectionCount() int {
+	return s.Stats.GetConnectionCount()
+}
+
 // NewGreeter creates new greeter server
 func NewGreeter() *Greeter {
 	streamData := []*HelloReply{
@@ -137,3 +151,50 @@ func NewGreeter() *Greeter {
 
 	return &Greeter{streamData: streamData, callCounts: m, mutex: &sync.RWMutex{}}
 }
+
+// NewHWStats creates new stats handler
+func NewHWStats() *HWStatsHandler {
+	return &HWStatsHandler{connCount: 0, mutex: &sync.RWMutex{}}
+}
+
+// HWStatsHandler is for gRPC stats
+type HWStatsHandler struct {
+	mutex     *sync.RWMutex
+	connCount int
+}
+
+// GetConnectionCount gets the connection count
+func (c *HWStatsHandler) GetConnectionCount() int {
+	c.mutex.RLock()
+	val := c.connCount
+	c.mutex.RUnlock()
+
+	return val
+}
+
+// HandleConn handle the connection
+func (c *HWStatsHandler) HandleConn(ctx context.Context, cs stats.ConnStats) {
+	fmt.Println("!!! HandleConn")
+	// no-op
+}
+
+// TagConn exists to satisfy gRPC stats.Handler.
+func (c *HWStatsHandler) TagConn(ctx context.Context, cti *stats.ConnTagInfo) context.Context {
+	fmt.Println("!!! TagConn")
+
+	c.mutex.Lock()
+	c.connCount++
+	c.mutex.Unlock()
+
+	return ctx
+}
+
+// HandleRPC implements per-RPC tracing and stats instrumentation.
+func (c *HWStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
+	// no-op
+}
+
+// TagRPC implements per-RPC context management.
+func (c *HWStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+	return ctx
+}
diff --git a/runner/options_test.go b/runner/options_test.go
@@ -278,6 +278,7 @@ func TestRunConfig_newRunConfig(t *testing.T) {
 	t.Run("with data from reader", func(t *testing.T) {
 
 		file, _ := os.Open("../testdata/data.json")
+		defer file.Close()
 
 		c, err := newConfig("call", "localhost:50050",
 			WithProtoFile("testdata/data.proto", []string{}),
@@ -292,6 +293,7 @@ func TestRunConfig_newRunConfig(t *testing.T) {
 		assert.Equal(t, 200, c.n)
 		assert.Equal(t, 50, c.c)
 		assert.Equal(t, 0, c.qps)
+		assert.Equal(t, 1, c.nConns)
 		assert.Equal(t, false, c.binary)
 		assert.Equal(t, time.Duration(0), c.z)
 		assert.Equal(t, time.Duration(0), c.keepaliveTime)
@@ -306,4 +308,55 @@ func TestRunConfig_newRunConfig(t *testing.T) {
 		assert.Equal(t, "", string(c.protoset))
 		assert.Equal(t, []string{"testdata", "."}, c.importPaths)
 	})
+
+	t.Run("with connections", func(t *testing.T) {
+
+		file, _ := os.Open("../testdata/data.json")
+		defer file.Close()
+
+		c, err := newConfig("call", "localhost:50050",
+			WithProtoFile("testdata/data.proto", []string{}),
+			WithDataFromReader(file),
+			WithConnections(5),
+		)
+
+		assert.NoError(t, err)
+
+		assert.Equal(t, "call", c.call)
+		assert.Equal(t, "localhost:50050", c.host)
+		assert.Equal(t, false, c.insecure)
+		assert.Equal(t, 200, c.n)
+		assert.Equal(t, 50, c.c)
+		assert.Equal(t, 0, c.qps)
+		assert.Equal(t, 5, c.nConns)
+		assert.Equal(t, false, c.binary)
+		assert.Equal(t, time.Duration(0), c.z)
+		assert.Equal(t, time.Duration(0), c.keepaliveTime)
+		assert.Equal(t, time.Duration(20*time.Second), c.timeout)
+		assert.Equal(t, time.Duration(10*time.Second), c.dialTimeout)
+		assert.Equal(t, runtime.GOMAXPROCS(-1), c.cpus)
+		assert.Empty(t, c.name)
+		assert.NotEmpty(t, c.data)
+		assert.False(t, c.binary)
+		assert.Empty(t, c.metadata)
+		assert.Equal(t, "testdata/data.proto", string(c.proto))
+		assert.Equal(t, "", string(c.protoset))
+		assert.Equal(t, []string{"testdata", "."}, c.importPaths)
+	})
+
+	t.Run("with invalid connections > concurrency", func(t *testing.T) {
+
+		file, _ := os.Open("../testdata/data.json")
+		defer file.Close()
+
+		c, err := newConfig("call", "localhost:50050",
+			WithProtoFile("testdata/data.proto", []string{}),
+			WithDataFromReader(file),
+			WithConcurrency(5),
+			WithConnections(6),
+		)
+
+		assert.Error(t, err)
+		assert.Nil(t, c)
+	})
 }
diff --git a/runner/run_test.go b/runner/run_test.go
@@ -108,6 +108,9 @@ func TestRunUnary(t *testing.T) {
 
 		count := gs.GetCount(callType)
 		assert.Equal(t, 12, count)
+
+		connCount := gs.GetConnectionCount()
+		assert.Equal(t, 1, connCount)
 	})
 
 	t.Run("test QPS", func(t *testing.T) {
@@ -210,6 +213,58 @@ func TestRunUnary(t *testing.T) {
 
 		count := gs.GetCount(callType)
 		assert.Equal(t, 5, count)
+
+		connCount := gs.GetConnectionCount()
+		assert.Equal(t, 1, connCount)
+	})
+
+	t.Run("test connections", func(t *testing.T) {
+		gs.ResetCounters()
+
+		msg := &helloworld.HelloRequest{}
+		msg.Name = "bob"
+
+		binData, err := proto.Marshal(msg)
+
+		report, err := Run(
+			"helloworld.Greeter.SayHello",
+			internal.TestLocalhost,
+			WithProtoFile("../testdata/greeter.proto", []string{}),
+			WithTotalRequests(5),
+			WithConcurrency(5),
+			WithConnections(5),
+			WithTimeout(time.Duration(20*time.Second)),
+			WithDialTimeout(time.Duration(20*time.Second)),
+			WithBinaryData(binData),
+			WithInsecure(true),
+		)
+
+		assert.NoError(t, err)
+
+		assert.NotNil(t, report)
+
+		assert.Equal(t, 5, int(report.Count))
+		assert.NotZero(t, report.Average)
+		assert.NotZero(t, report.Fastest)
+		assert.NotZero(t, report.Slowest)
+		assert.NotZero(t, report.Rps)
+		assert.Empty(t, report.Name)
+		assert.NotEmpty(t, report.Date)
+		assert.NotEmpty(t, report.Options)
+		assert.NotEmpty(t, report.Details)
+		assert.NotEmpty(t, report.LatencyDistribution)
+		assert.Equal(t, ReasonNormalEnd, report.EndReason)
+		assert.Empty(t, report.ErrorDist)
+
+		assert.NotEqual(t, report.Average, report.Slowest)
+		assert.NotEqual(t, report.Average, report.Fastest)
+		assert.NotEqual(t, report.Slowest, report.Fastest)
+
+		count := gs.GetCount(callType)
+		assert.Equal(t, 5, count)
+
+		connCount := gs.GetConnectionCount()
+		assert.Equal(t, 5, connCount)
 	})
 }
 
@@ -266,6 +321,9 @@ func TestRunServerStreaming(t *testing.T) {
 
 	count := gs.GetCount(callType)
 	assert.Equal(t, 15, count)
+
+	connCount := gs.GetConnectionCount()
+	assert.Equal(t, 1, connCount)
 }
 
 func TestRunClientStreaming(t *testing.T) {
@@ -328,6 +386,9 @@ func TestRunClientStreaming(t *testing.T) {
 
 	count := gs.GetCount(callType)
 	assert.Equal(t, 16, count)
+
+	connCount := gs.GetConnectionCount()
+	assert.Equal(t, 1, connCount)
 }
 
 func TestRunClientStreamingBinary(t *testing.T) {
@@ -384,6 +445,9 @@ func TestRunClientStreamingBinary(t *testing.T) {
 
 	count := gs.GetCount(callType)
 	assert.Equal(t, 24, count)
+
+	connCount := gs.GetConnectionCount()
+	assert.Equal(t, 1, connCount)
 }
 
 func TestRunBidi(t *testing.T) {
@@ -446,6 +510,9 @@ func TestRunBidi(t *testing.T) {
 
 	count := gs.GetCount(callType)
 	assert.Equal(t, 20, count)
+
+	connCount := gs.GetConnectionCount()
+	assert.Equal(t, 1, connCount)
 }
 
 func TestRunUnarySecure(t *testing.T) {
@@ -500,6 +567,9 @@ func TestRunUnarySecure(t *testing.T) {
 
 	count := gs.GetCount(callType)
 	assert.Equal(t, 18, count)
+
+	connCount := gs.GetConnectionCount()
+	assert.Equal(t, 1, connCount)
 }
 
 func TestRunUnaryProtoset(t *testing.T) {
@@ -560,6 +630,9 @@ func TestRunUnaryProtoset(t *testing.T) {
 
 	count := gs.GetCount(callType)
 	assert.Equal(t, 21, count)
+
+	connCount := gs.GetConnectionCount()
+	assert.Equal(t, 1, connCount)
 }
 
 func TestRunUnaryReflection(t *testing.T) {
@@ -594,6 +667,9 @@ func TestRunUnaryReflection(t *testing.T) {
 
 		assert.Error(t, err)
 		assert.Nil(t, report)
+
+		connCount := gs.GetConnectionCount()
+		assert.Equal(t, 1, connCount)
 	})
 
 	t.Run("Unary streaming", func(t *testing.T) {
@@ -656,6 +732,9 @@ func TestRunUnaryReflection(t *testing.T) {
 
 		count := gs.GetCount(callType)
 		assert.Equal(t, 21, count)
+
+		connCount := gs.GetConnectionCount()
+		assert.Equal(t, 2, connCount) // 1 extra connection for reflection
 	})
 
 	t.Run("Client streaming", func(t *testing.T) {
@@ -717,5 +796,8 @@ func TestRunUnaryReflection(t *testing.T) {
 
 		count := gs.GetCount(callType)
 		assert.Equal(t, 16, count)
+
+		connCount := gs.GetConnectionCount()
+		assert.Equal(t, 2, connCount) // 1 extra connection for reflection
 	})
 }