fix: startWithListener register etcd

zeromicro · Dec 17, 2024 · 022b53b · 022b53b
1 parent 26450b7
commit 022b53b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/zrpc/internal/rpcpubserver.go b/zrpc/internal/rpcpubserver.go
@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"net"
 	"os"
 	"strings"
 
@@ -53,6 +54,14 @@ func (s keepAliveServer) Start(fn RegisterFn) error {
 	return s.Server.Start(fn)
 }
 
+func (s keepAliveServer) StartWithListener(listener net.Listener, fn RegisterFn) error {
+	if err := s.registerEtcd(); err != nil {
+		return err
+	}
+
+	return s.Server.StartWithListener(listener, fn)
+}
+
 func figureOutListenOn(listenOn string) string {
 	fields := strings.Split(listenOn, ":")
 	if len(fields) == 0 {

diff --git a/zrpc/internal/rpcpubserver_test.go b/zrpc/internal/rpcpubserver_test.go
@@ -18,6 +18,10 @@ func TestNewRpcPubServer(t *testing.T) {
 	assert.NotPanics(t, func() {
 		s.Start(nil)
 	})
+
+	assert.NotPanics(t, func() {
+		s.StartWithListener(nil, nil)
+	})
 }
 
 func TestFigureOutListenOn(t *testing.T) {