Decouple registration from grpc.Server implementation

server.go

@@ -30,11 +30,15 @@ func init() {
 	prom.MustRegister(DefaultServerMetrics.serverStreamMsgSent)
 }
 
-// Register takes a gRPC server 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) {
+// prometheusServer defines the interface that is required for grpc_prometheus to Register.
+type prometheusServer interface {
+	GetServiceInfo() map[string]grpc.ServiceInfo
+}
+
+// Register takes a prometheusServer 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.
+func Register(server prometheusServer) {
 	DefaultServerMetrics.InitializeMetrics(server)
 }
 

server_metrics.go

@@ -129,9 +129,9 @@ 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
+// value, for all gRPC methods registered on a prometheusServer. This is useful, to
 // ensure that all metrics exist when collecting and querying.
-func (m *ServerMetrics) InitializeMetrics(server *grpc.Server) {
+func (m *ServerMetrics) InitializeMetrics(server prometheusServer) {
 	serviceInfo := server.GetServiceInfo()
 	for serviceName, info := range serviceInfo {
 		for _, mInfo := range info.Methods {