gorush/rpc/server_test.go

package rpc

import (
	"context"
	"testing"

	"github.com/appleboy/gorush/config"
	"github.com/appleboy/gorush/logx"

	"google.golang.org/grpc"
	"google.golang.org/grpc/connectivity"
)

const gRPCAddr = "localhost:9000"

func initTest() config.ConfYaml {
	cfg, _ := config.LoadConf()
	cfg.Core.Mode = "test"
	return cfg
}

func TestGracefulShutDownGRPCServer(t *testing.T) {
	cfg := initTest()
	// server configs
	logx.InitLog(
		cfg.Log.AccessLevel,
		cfg.Log.AccessLog,
		cfg.Log.ErrorLevel,
		cfg.Log.ErrorLog,
	)
	cfg.GRPC.Enabled = true
	cfg.GRPC.Port = "9000"
	cfg.Log.Format = "json"

	// Run gRPC server
	ctx, gRPCContextCancel := context.WithCancel(context.Background())
	go func() {
		if err := RunGRPCServer(ctx, cfg); err != nil {
			panic(err)
		}
	}()

	// gRPC client conn
	conn, err := grpc.Dial(
		gRPCAddr,
		grpc.WithInsecure(),
		grpc.WithDefaultCallOptions(grpc.WaitForReady(true)),
	) // wait for server ready
	if err != nil {
		t.Error(err)
	}

	// Stop gRPC server
	go gRPCContextCancel()

	// wait for client connection would be closed
	for conn.GetState() != connectivity.TransientFailure {
	}
	conn.Close()
}