use alternate interface; change signature of Register

reflection/serverreflection.go

@@ -54,33 +54,29 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+// GRPCServer is the interface provided by a gRPC server. It is implemented by
+// *grpc.Server, but could also be implemented by other concrete types. It acts
+// as a registry, for accumulating the services exposed by the server.
+type GRPCServer interface {
+	grpc.ServiceRegistrar
+	GetServiceInfo() map[string]grpc.ServiceInfo
+}
+
+var _ GRPCServer = (*grpc.Server)(nil)
+
 type serverReflectionServer struct {
 	rpb.UnimplementedServerReflectionServer
-	queryServices func() map[string]grpc.ServiceInfo
+	s GRPCServer
 
 	initSymbols  sync.Once
 	serviceNames []string
 	symbols      map[string]*dpb.FileDescriptorProto // map of fully-qualified names to files
 }
 
 // Register registers the server reflection service on the given gRPC server.
-func Register(s *grpc.Server) {
-	register(s, s.GetServiceInfo)
-}
-
-// RegisterWithRegistrar registers the server reflection service with the given registrar.
-// The given map describes the services that will be exposed via reflection.
-func RegisterWithRegistrar(r grpc.ServiceRegistrar, svcs map[string]grpc.ServiceInfo) {
-	rpb.RegisterServerReflectionServer(r, &serverReflectionServer{
-		queryServices: func() map[string]grpc.ServiceInfo {
-			return svcs
-		},
-	})
-}
-
-func register(r grpc.ServiceRegistrar, queryServices func() map[string]grpc.ServiceInfo) {
-	rpb.RegisterServerReflectionServer(r, &serverReflectionServer{
-		queryServices: queryServices,
+func Register(s GRPCServer) {
+	rpb.RegisterServerReflectionServer(s, &serverReflectionServer{
+		s: s,
 	})
 }
 
@@ -94,7 +90,7 @@ type protoMessage interface {
 
 func (s *serverReflectionServer) getSymbols() (svcNames []string, symbolIndex map[string]*dpb.FileDescriptorProto) {
 	s.initSymbols.Do(func() {
-		serviceInfo := s.queryServices()
+		serviceInfo := s.s.GetServiceInfo()
 
 		s.symbols = map[string]*dpb.FileDescriptorProto{}
 		s.serviceNames = make([]string, 0, len(serviceInfo))