Fix bugs in extension treatment REST API

Change-Id: If47f5ffce634617200e765be72c7898d984b0786
diff --git a/core/common/src/main/java/org/onosproject/codec/impl/DecodeInstructionCodecHelper.java b/core/common/src/main/java/org/onosproject/codec/impl/DecodeInstructionCodecHelper.java
index b70a684..5e69b0e 100644
--- a/core/common/src/main/java/org/onosproject/codec/impl/DecodeInstructionCodecHelper.java
+++ b/core/common/src/main/java/org/onosproject/codec/impl/DecodeInstructionCodecHelper.java
@@ -25,6 +25,7 @@
 import org.onlab.packet.TpPort;
 import org.onlab.packet.VlanId;
 import org.onlab.util.HexString;
+import org.onosproject.codec.CodecContext;
 import org.onosproject.codec.ExtensionTreatmentCodec;
 import org.onosproject.core.DefaultGroupId;
 import org.onosproject.core.GroupId;
@@ -56,14 +57,16 @@
 public final class DecodeInstructionCodecHelper {
     protected static final Logger log = getLogger(DecodeInstructionCodecHelper.class);
     private final ObjectNode json;
+    private final CodecContext context;
 
     /**
      * Creates a decode instruction codec object.
      *
      * @param json JSON object to decode
      */
-    public DecodeInstructionCodecHelper(ObjectNode json) {
+    public DecodeInstructionCodecHelper(ObjectNode json, CodecContext context) {
         this.json = json;
+        this.context = context;
     }
 
     /**
@@ -251,12 +254,17 @@
             DeviceService deviceService = serviceDirectory.get(DeviceService.class);
             Device device = deviceService.getDevice(deviceId);
 
+            if (device == null) {
+                throw new IllegalArgumentException("Device not found");
+            }
+
             if (device.is(ExtensionTreatmentCodec.class)) {
                 ExtensionTreatmentCodec treatmentCodec = device.as(ExtensionTreatmentCodec.class);
-                ExtensionTreatment treatment = treatmentCodec.decode(node, null);
+                ExtensionTreatment treatment = treatmentCodec.decode(node, context);
                 return Instructions.extension(treatment, deviceId);
             } else {
-                log.warn("There is no codec to decode extension for device {}", deviceId.toString());
+                throw new IllegalArgumentException(
+                        "There is no codec to decode extension for device " + deviceId.toString());
             }
         }
         return null;
diff --git a/core/common/src/main/java/org/onosproject/codec/impl/EncodeInstructionCodecHelper.java b/core/common/src/main/java/org/onosproject/codec/impl/EncodeInstructionCodecHelper.java
index d04acf5..2b3c9d7 100644
--- a/core/common/src/main/java/org/onosproject/codec/impl/EncodeInstructionCodecHelper.java
+++ b/core/common/src/main/java/org/onosproject/codec/impl/EncodeInstructionCodecHelper.java
@@ -241,12 +241,17 @@
         DeviceService deviceService = serviceDirectory.get(DeviceService.class);
         Device device = deviceService.getDevice(deviceId);
 
+        if (device == null) {
+            throw new IllegalArgumentException("Device not found");
+        }
+
         if (device.is(ExtensionTreatmentCodec.class)) {
             ExtensionTreatmentCodec treatmentCodec = device.as(ExtensionTreatmentCodec.class);
             ObjectNode node = treatmentCodec.encode(extensionInstruction.extensionInstruction(), context);
             result.set(InstructionCodec.EXTENSION, node);
         } else {
-            log.warn("There is no codec to encode extension for device {}", deviceId.toString());
+            throw new IllegalArgumentException(
+                    "There is no codec to encode extension for device " + deviceId.toString());
         }
     }
 
diff --git a/core/common/src/main/java/org/onosproject/codec/impl/InstructionCodec.java b/core/common/src/main/java/org/onosproject/codec/impl/InstructionCodec.java
index da41e25..9da240d 100644
--- a/core/common/src/main/java/org/onosproject/codec/impl/InstructionCodec.java
+++ b/core/common/src/main/java/org/onosproject/codec/impl/InstructionCodec.java
@@ -77,6 +77,6 @@
             return null;
         }
 
-        return new DecodeInstructionCodecHelper(json).decode();
+        return new DecodeInstructionCodecHelper(json, context).decode();
     }
 }
diff --git a/drivers/default/src/main/resources/onos-drivers.xml b/drivers/default/src/main/resources/onos-drivers.xml
index 673befa..3daf2e3 100644
--- a/drivers/default/src/main/resources/onos-drivers.xml
+++ b/drivers/default/src/main/resources/onos-drivers.xml
@@ -30,10 +30,14 @@
                    impl="org.onosproject.driver.extensions.NiciraExtensionTreatmentInterpreter" />
         <behaviour api="org.onosproject.net.behaviour.ExtensionTreatmentResolver"
                    impl="org.onosproject.driver.extensions.NiciraExtensionTreatmentInterpreter" />
+        <behaviour api="org.onosproject.codec.ExtensionTreatmentCodec"
+                   impl="org.onosproject.driver.extensions.NiciraExtensionTreatmentInterpreter" />
         <behaviour api="org.onosproject.openflow.controller.ExtensionSelectorInterpreter"
                    impl="org.onosproject.driver.extensions.NiciraExtensionSelectorInterpreter" />
         <behaviour api="org.onosproject.net.behaviour.ExtensionSelectorResolver"
                    impl="org.onosproject.driver.extensions.NiciraExtensionSelectorInterpreter" />
+        <behaviour api="org.onosproject.codec.ExtensionSelectorCodec"
+                   impl="org.onosproject.driver.extensions.NiciraExtensionSelectorInterpreter" />
         <behaviour api="org.onosproject.net.behaviour.VlanQuery"
                    impl="org.onosproject.driver.query.FullVlanAvailable" />
         <behaviour api="org.onosproject.net.behaviour.MplsQuery"