diff --git a/drivers/gnmi/src/main/java/org/onosproject/drivers/gnmi/OpenConfigGnmiPortStatisticsDiscovery.java b/drivers/gnmi/src/main/java/org/onosproject/drivers/gnmi/OpenConfigGnmiPortStatisticsDiscovery.java
index fa7231f..6e361ff 100644
--- a/drivers/gnmi/src/main/java/org/onosproject/drivers/gnmi/OpenConfigGnmiPortStatisticsDiscovery.java
+++ b/drivers/gnmi/src/main/java/org/onosproject/drivers/gnmi/OpenConfigGnmiPortStatisticsDiscovery.java
@@ -23,6 +23,9 @@
 import gnmi.Gnmi.GetResponse;
 import gnmi.Gnmi.Path;
 import org.apache.commons.lang3.tuple.Pair;
+import org.onosproject.gnmi.api.GnmiClient;
+import org.onosproject.gnmi.api.GnmiController;
+import org.onosproject.grpc.utils.AbstractGrpcHandlerBehaviour;
 import org.onosproject.net.DeviceId;
 import org.onosproject.net.Port;
 import org.onosproject.net.PortNumber;
@@ -41,13 +44,18 @@
 /**
  * Behaviour to get port statistics from device via gNMI.
  */
-public class OpenConfigGnmiPortStatisticsDiscovery extends AbstractGnmiHandlerBehaviour
+public class OpenConfigGnmiPortStatisticsDiscovery
+        extends AbstractGrpcHandlerBehaviour<GnmiClient, GnmiController>
         implements PortStatisticsDiscovery {
 
     private static final Map<Pair<DeviceId, PortNumber>, Long> PORT_START_TIMES =
             Maps.newConcurrentMap();
     private static final String LAST_CHANGE = "last-changed";
 
+    public OpenConfigGnmiPortStatisticsDiscovery() {
+        super(GnmiController.class);
+    }
+
     @Override
     public Collection<PortStatistics> discoverPortStatistics() {
         if (!setupBehaviour("discoverPortStatistics()")) {