From 9b114e3251b2ce583eb52fa299cb16bf608cfab4 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Fri, 8 Oct 2021 17:31:37 +0800
Subject: [PATCH] test: add more tests (#1119)

---
 zrpc/internal/rpcserver_test.go | 39 ++++++++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/zrpc/internal/rpcserver_test.go b/zrpc/internal/rpcserver_test.go
index 9c1a9b9e0..76e55d99c 100644
--- a/zrpc/internal/rpcserver_test.go
+++ b/zrpc/internal/rpcserver_test.go
@@ -1,16 +1,45 @@
 package internal
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/stat"
+	"github.com/tal-tech/go-zero/zrpc/internal/mock"
+	"google.golang.org/grpc"
 )
 
-func TestWithMetrics(t *testing.T) {
+func TestRpcServer(t *testing.T) {
 	metrics := stat.NewMetrics("foo")
-	opt := WithMetrics(metrics)
-	var options rpcServerOptions
-	opt(&options)
-	assert.Equal(t, metrics, options.metrics)
+	server := NewRpcServer("localhost:54321", WithMetrics(metrics))
+	server.SetName("mock")
+	var wg sync.WaitGroup
+	var grpcServer *grpc.Server
+	var lock sync.Mutex
+	wg.Add(1)
+	go func() {
+		err := server.Start(func(server *grpc.Server) {
+			lock.Lock()
+			mock.RegisterDepositServiceServer(server, new(mock.DepositServer))
+			grpcServer = server
+			lock.Unlock()
+			wg.Done()
+		})
+		assert.Nil(t, err)
+	}()
+
+	wg.Wait()
+	lock.Lock()
+	grpcServer.GracefulStop()
+	lock.Unlock()
+}
+
+func TestRpcServer_WithBadAddress(t *testing.T) {
+	server := NewRpcServer("localhost:111111")
+	server.SetName("mock")
+	err := server.Start(func(server *grpc.Server) {
+		mock.RegisterDepositServiceServer(server, new(mock.DepositServer))
+	})
+	assert.NotNil(t, err)
 }