From 25261719a02bbe29263bbebac6a6c5b25f117f68 Mon Sep 17 00:00:00 2001
From: Tony Tkacik <ttkacik@cisco.com>
Date: Mon, 1 Sep 2014 21:57:47 +0200
Subject: [PATCH] Fixed SchemaTracker behaviour in choice nodes.

Change-Id: Ia1ed56a8a263b1d0397ad0c54f9b6e5775939d89
Signed-off-by: Tony Tkacik <ttkacik@cisco.com>
---
 .../gson/JSONNormalizedNodeStreamWriter.java  | 31 +++++++++++++++++--
 .../yang/data/impl/codec/SchemaTracker.java   | 19 ++++++++++--
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/yang/yang-data-codec-gson/src/main/java/org/opendaylight/yangtools/yang/data/codec/gson/JSONNormalizedNodeStreamWriter.java b/yang/yang-data-codec-gson/src/main/java/org/opendaylight/yangtools/yang/data/codec/gson/JSONNormalizedNodeStreamWriter.java
index aa4936a595..e26570eb30 100644
--- a/yang/yang-data-codec-gson/src/main/java/org/opendaylight/yangtools/yang/data/codec/gson/JSONNormalizedNodeStreamWriter.java
+++ b/yang/yang-data-codec-gson/src/main/java/org/opendaylight/yangtools/yang/data/codec/gson/JSONNormalizedNodeStreamWriter.java
@@ -10,13 +10,11 @@ package org.opendaylight.yangtools.yang.data.codec.gson;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.gson.stream.JsonWriter;
-
 import java.io.IOException;
 import java.io.Writer;
 import java.net.URI;
 import java.util.ArrayDeque;
 import java.util.Deque;
-
 import org.opendaylight.yangtools.concepts.Codec;
 import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.AugmentationIdentifier;
 import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.NodeIdentifier;
@@ -30,6 +28,7 @@ import org.opendaylight.yangtools.yang.model.api.LeafListSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.LeafSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.Module;
 import org.opendaylight.yangtools.yang.model.api.SchemaContext;
+import org.opendaylight.yangtools.yang.model.api.SchemaPath;
 
 /**
  * This implementation will create JSON output as output stream.
@@ -101,6 +100,23 @@ public class JSONNormalizedNodeStreamWriter implements NormalizedNodeStreamWrite
         this.tracker = SchemaTracker.create(schemaContext);
     }
 
+    private JSONNormalizedNodeStreamWriter(final SchemaContext schemaContext, final SchemaPath path,
+            final Writer writer, final int indentSize) {
+        this.schemaContext = Preconditions.checkNotNull(schemaContext);
+        this.writer = Preconditions.checkNotNull(writer);
+
+        Preconditions.checkArgument(indentSize >= 0, "Indent size must be non-negative");
+        if (indentSize != 0) {
+            indent = Strings.repeat(" ", indentSize);
+        } else {
+            indent = null;
+        }
+
+        this.utils = SchemaContextUtils.create(schemaContext);
+        this.codecs = RestCodecFactory.create(utils);
+        this.tracker = SchemaTracker.create(schemaContext,path);
+    }
+
     /**
      * Create a new stream writer, which writes to the specified {@link Writer}.
      *
@@ -112,6 +128,17 @@ public class JSONNormalizedNodeStreamWriter implements NormalizedNodeStreamWrite
         return new JSONNormalizedNodeStreamWriter(schemaContext, writer, 0);
     }
 
+    /**
+     * Create a new stream writer, which writes to the specified {@link Writer}.
+     *
+     * @param schemaContext Schema context
+     * @param writer Output writer
+     * @return A stream writer instance
+     */
+    public static NormalizedNodeStreamWriter create(final SchemaContext schemaContext, SchemaPath path,final Writer writer) {
+        return new JSONNormalizedNodeStreamWriter(schemaContext, path, writer, 0);
+    }
+
     /**
      * Create a new stream writer, which writes to the specified output stream.
      *
diff --git a/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java b/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java
index 0a41e28f93..4a8232e180 100644
--- a/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java
+++ b/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java
@@ -24,6 +24,7 @@ import org.opendaylight.yangtools.yang.data.impl.schema.transform.base.Augmentat
 import org.opendaylight.yangtools.yang.model.api.AnyXmlSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.AugmentationSchema;
 import org.opendaylight.yangtools.yang.model.api.AugmentationTarget;
+import org.opendaylight.yangtools.yang.model.api.ChoiceCaseNode;
 import org.opendaylight.yangtools.yang.model.api.ChoiceNode;
 import org.opendaylight.yangtools.yang.model.api.ContainerSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.DataNodeContainer;
@@ -90,10 +91,22 @@ public final class SchemaTracker {
 
     private SchemaNode getSchema(final PathArgument name) {
         final Object parent = getParent();
-        Preconditions.checkState(parent instanceof DataNodeContainer);
-
+        SchemaNode schema = null;
         final QName qname = name.getNodeType();
-        final SchemaNode schema = ((DataNodeContainer)parent).getDataChildByName(qname);
+        if(parent instanceof DataNodeContainer) {
+            schema = ((DataNodeContainer)parent).getDataChildByName(qname);
+
+        } else if(parent instanceof ChoiceNode) {
+            for(ChoiceCaseNode caze : ((ChoiceNode) parent).getCases()) {
+                DataSchemaNode potential = caze.getDataChildByName(qname);
+                if(potential != null) {
+                    schema = potential;
+                    break;
+                }
+            }
+        } else {
+            throw new IllegalStateException("Unsupported schema type "+ parent.getClass() +" on stack.");
+        }
         Preconditions.checkArgument(schema != null, "Could not find schema for node %s in %s", qname, parent);
         return schema;
     }
-- 
2.36.6