diff --git a/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/AbstractGnoiHandlerBehaviour.java b/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/AbstractGnoiHandlerBehaviour.java
index 23181b4..84043a8 100644
--- a/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/AbstractGnoiHandlerBehaviour.java
+++ b/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/AbstractGnoiHandlerBehaviour.java
@@ -20,11 +20,9 @@
 import org.onosproject.gnoi.api.GnoiClient;
 import org.onosproject.gnoi.api.GnoiClientKey;
 import org.onosproject.gnoi.api.GnoiController;
-import org.onosproject.net.Device;
 import org.onosproject.net.DeviceId;
 import org.onosproject.net.config.NetworkConfigService;
 import org.onosproject.net.config.basics.BasicDeviceConfig;
-import org.onosproject.net.device.DeviceService;
 import org.onosproject.net.driver.AbstractHandlerBehaviour;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -39,19 +37,14 @@
 
     protected final Logger log = LoggerFactory.getLogger(getClass());
     protected DeviceId deviceId;
-    protected DeviceService deviceService;
-    protected Device device;
-    protected GnoiController controller;
     protected GnoiClient client;
 
-    protected boolean setupBehaviour() {
+    protected boolean setupBehaviour(String opName) {
         deviceId = handler().data().deviceId();
-        deviceService = handler().get(DeviceService.class);
-        controller = handler().get(GnoiController.class);
-        client = controller.getClient(deviceId);
+        client = getClientByKey();
 
         if (client == null) {
-            log.warn("Unable to find client for {}, aborting operation", deviceId);
+            log.warn("Missing client for {}, aborting {}", deviceId, opName);
             return false;
         }
 
diff --git a/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/GnoiBasicSystemOperationsImpl.java b/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/GnoiBasicSystemOperationsImpl.java
index 411a8bf..b979546 100644
--- a/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/GnoiBasicSystemOperationsImpl.java
+++ b/drivers/gnoi/src/main/java/org/onosproject/drivers/gnoi/GnoiBasicSystemOperationsImpl.java
@@ -16,17 +16,18 @@
 
 package org.onosproject.drivers.gnoi;
 
+import gnoi.system.SystemOuterClass.RebootMethod;
 import gnoi.system.SystemOuterClass.RebootRequest;
 import gnoi.system.SystemOuterClass.RebootResponse;
-import gnoi.system.SystemOuterClass.RebootMethod;
 import org.onosproject.net.behaviour.BasicSystemOperations;
-import java.util.concurrent.CompletableFuture;
-
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.concurrent.CompletableFuture;
+
 /**
- * Implementation of the BasicSystemOperations behavior for gNOI-enabled devices.
+ * Implementation of the BasicSystemOperations behavior for gNOI-enabled
+ * devices.
  */
 public class GnoiBasicSystemOperationsImpl
         extends AbstractGnoiHandlerBehaviour implements BasicSystemOperations {
@@ -36,12 +37,13 @@
 
     @Override
     public CompletableFuture<Boolean> reboot() {
-        if (!setupBehaviour()) {
+        if (!setupBehaviour("reboot()")) {
             return CompletableFuture.completedFuture(false);
         }
 
         final RebootRequest.Builder requestMsg = RebootRequest.newBuilder().setMethod(RebootMethod.COLD);
-        final CompletableFuture<Boolean> future = client.reboot(requestMsg.build())
+
+        return client.reboot(requestMsg.build())
                 .handle((response, error) -> {
                     if (error == null) {
                         log.debug("gNOI reboot() for device {} returned {}", deviceId, response);
@@ -51,16 +53,14 @@
                         return false;
                     }
                 });
-
-        return future;
     }
 
     @Override
     public CompletableFuture<Long> time() {
-        if (!setupBehaviour()) {
+        if (!setupBehaviour("time()")) {
             return CompletableFuture.completedFuture(0L);
         }
-        final CompletableFuture<Long> future = client.time()
+        return client.time()
                 .handle((response, error) -> {
                     if (error == null) {
                         log.debug("gNOI time() for device {} returned {}", deviceId.uri(), response.getTime());
@@ -70,7 +70,5 @@
                         return 0L;
                     }
                 });
-
-        return future;
     }
 }