Decouple registration from grpc.Server implementation

server.go

@@ -30,11 +30,16 @@ func init() {
 	prom.MustRegister(DefaultServerMetrics.serverStreamMsgSent)
 }
 
-// Register takes a gRPC server and pre-initializes all counters to 0. This
+// prometheusServer defines the interface that is required for grpc_prometheus to Register.
+type prometheusServer interface {
+	GetServiceInfo() map[string]grpc.ServiceInfo
+}
+
+// Register takes a gRPC server (or any other struct that meets the
+// prometheusServer interface) and pre-initializes all counters to 0. This
 // allows for easier monitoring in Prometheus (no missing metrics), and should
-// be called *after* all services have been registered with the server. This
-// function acts on the DefaultServerMetrics variable.
-func Register(server *grpc.Server) {
+// be called *after* all services have been registered with the server.
+func Register(server prometheusServer) {
 	DefaultServerMetrics.InitializeMetrics(server)
 }
 

server_metrics.go

@@ -126,9 +126,10 @@ func (m *ServerMetrics) StreamServerInterceptor() func(srv interface{}, ss grpc.
 }
 
 // InitializeMetrics initializes all metrics, with their appropriate null
-// value, for all gRPC methods registered on a gRPC server. This is useful, to
-// ensure that all metrics exist when collecting and querying.
-func (m *ServerMetrics) InitializeMetrics(server *grpc.Server) {
+// value, for all gRPC methods registered on a gRPC server (or any other
+// struct that meets the prometheusServer interface). This is useful, to ensure
+// that all metrics exist when collecting and querying.
+func (m *ServerMetrics) InitializeMetrics(server prometheusServer) {
 	serviceInfo := server.GetServiceInfo()
 	for serviceName, info := range serviceInfo {
 		for _, mInfo := range info.Methods {