BUG-1513: migrate users to ChoiceSchemaNode
[controller.git] / opendaylight / md-sal / sal-rest-docgen / src / main / java / org / opendaylight / controller / sal / rest / doc / impl / ModelGenerator.java
index 719dd78064fbee1815078866de80f2f90809ef85..65c835723df3fcf31a5c1b57d9f04d77ebf93fab 100644 (file)
@@ -7,6 +7,11 @@
  */
 package org.opendaylight.controller.sal.rest.doc.impl;
 
+import static org.opendaylight.controller.sal.rest.doc.impl.BaseYangSwaggerGenerator.MODULE_NAME_SUFFIX;
+import static org.opendaylight.controller.sal.rest.doc.model.builder.OperationBuilder.Post.METHOD_NAME;
+import static org.opendaylight.controller.sal.rest.doc.util.RestDocgenUtil.resolveNodesName;
+
+import com.google.common.base.Preconditions;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -14,21 +19,27 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-
+import org.apache.commons.lang3.BooleanUtils;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
+import org.opendaylight.controller.sal.rest.doc.model.builder.OperationBuilder;
+import org.opendaylight.yangtools.yang.common.QName;
 import org.opendaylight.yangtools.yang.model.api.AnyXmlSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.ChoiceCaseNode;
-import org.opendaylight.yangtools.yang.model.api.ChoiceNode;
+import org.opendaylight.yangtools.yang.model.api.ChoiceSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.ConstraintDefinition;
 import org.opendaylight.yangtools.yang.model.api.ContainerSchemaNode;
+import org.opendaylight.yangtools.yang.model.api.DataNodeContainer;
 import org.opendaylight.yangtools.yang.model.api.DataSchemaNode;
+import org.opendaylight.yangtools.yang.model.api.IdentitySchemaNode;
 import org.opendaylight.yangtools.yang.model.api.LeafListSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.LeafSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.ListSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.Module;
 import org.opendaylight.yangtools.yang.model.api.RpcDefinition;
+import org.opendaylight.yangtools.yang.model.api.SchemaContext;
+import org.opendaylight.yangtools.yang.model.api.SchemaNode;
 import org.opendaylight.yangtools.yang.model.api.TypeDefinition;
 import org.opendaylight.yangtools.yang.model.api.type.BinaryTypeDefinition;
 import org.opendaylight.yangtools.yang.model.api.type.BitsTypeDefinition;
@@ -58,7 +69,7 @@ import org.slf4j.LoggerFactory;
  */
 public class ModelGenerator {
 
-    private static final Logger _logger = LoggerFactory.getLogger(ModelGenerator.class);
+    private static final Logger LOG = LoggerFactory.getLogger(ModelGenerator.class);
 
     private static final String BASE_64 = "base64";
     private static final String BINARY_ENCODING_KEY = "binaryEncoding";
@@ -84,166 +95,279 @@ public class ModelGenerator {
     private static final String NUMBER = "number";
     private static final String BOOLEAN = "boolean";
     private static final String STRING = "string";
+    private static final String ID_KEY = "id";
+    private static final String SUB_TYPES_KEY = "subTypes";
 
     private static final Map<Class<? extends TypeDefinition<?>>, String> YANG_TYPE_TO_JSON_TYPE_MAPPING;
 
     static {
-        Map<Class<? extends TypeDefinition<?>>, String> tempMap1 = new HashMap<Class<? extends TypeDefinition<?>>, String>(10);
-        tempMap1.put(StringType.class , STRING);
-        tempMap1.put(BooleanType.class , BOOLEAN);
-        tempMap1.put(Int8.class , INTEGER);
-        tempMap1.put(Int16.class , INTEGER);
-        tempMap1.put(Int32.class , INTEGER);
-        tempMap1.put(Int64.class , INTEGER);
-        tempMap1.put(Uint16.class , INTEGER);
-        tempMap1.put(Uint32.class , INTEGER);
-        tempMap1.put(Uint64.class , INTEGER);
-        tempMap1.put(Uint8.class , INTEGER);
-        tempMap1.put(Decimal64.class , NUMBER);
-        tempMap1.put(EnumerationType.class , ENUM);
-        //TODO: Binary type
+        Map<Class<? extends TypeDefinition<?>>, String> tempMap1 = new HashMap<Class<? extends TypeDefinition<?>>, String>(
+                10);
+        tempMap1.put(StringType.class, STRING);
+        tempMap1.put(BooleanType.class, BOOLEAN);
+        tempMap1.put(Int8.class, INTEGER);
+        tempMap1.put(Int16.class, INTEGER);
+        tempMap1.put(Int32.class, INTEGER);
+        tempMap1.put(Int64.class, INTEGER);
+        tempMap1.put(Uint16.class, INTEGER);
+        tempMap1.put(Uint32.class, INTEGER);
+        tempMap1.put(Uint64.class, INTEGER);
+        tempMap1.put(Uint8.class, INTEGER);
+        tempMap1.put(Decimal64.class, NUMBER);
+        tempMap1.put(EnumerationType.class, ENUM);
+        // TODO: Binary type
 
         YANG_TYPE_TO_JSON_TYPE_MAPPING = Collections.unmodifiableMap(tempMap1);
     }
 
-    public ModelGenerator(){
+    private Module topLevelModule;
+
+    public ModelGenerator() {
     }
 
-    public JSONObject convertToJsonSchema(final Module module) throws IOException, JSONException {
+    public JSONObject convertToJsonSchema(Module module, SchemaContext schemaContext) throws IOException, JSONException {
         JSONObject models = new JSONObject();
-        processContainers(module, models);
-        processRPCs(module, models);
-
+        topLevelModule = module;
+        processModules(module, models);
+        processContainersAndLists(module, models, schemaContext);
+        processRPCs(module, models, schemaContext);
+        processIdentities(module, models);
         return models;
     }
 
+    private void processModules(Module module, JSONObject models) throws JSONException {
+        createConcreteModelForPost(models, module.getName()+MODULE_NAME_SUFFIX, createPropertiesForPost(module));
+    }
 
-
-    private void processContainers(final Module module, final JSONObject models) throws IOException, JSONException {
+    private void processContainersAndLists(Module module, JSONObject models, SchemaContext schemaContext)
+            throws IOException, JSONException {
 
         String moduleName = module.getName();
-        Set<DataSchemaNode> childNodes =  module.getChildNodes();
-
-        for(DataSchemaNode childNode : childNodes){
-            JSONObject moduleJSON=null;
-            String filename = childNode.getQName().getLocalName();
-            /*
-             * For every container in the module
-             */
-            if(childNode instanceof ContainerSchemaNode) {
-                moduleJSON = processContainer((ContainerSchemaNode)childNode, moduleName, true, models);
-            }
 
-            if(moduleJSON!=null) {
-                _logger.debug("Adding model for [{}]", filename);
-                moduleJSON.put("id", filename);
-                models.put(filename, moduleJSON);
+        for (DataSchemaNode childNode : module.getChildNodes()) {
+            // For every container and list in the module
+            if (childNode instanceof ContainerSchemaNode || childNode instanceof ListSchemaNode) {
+                processDataNodeContainer((DataNodeContainer) childNode, moduleName, models, true, schemaContext);
+                processDataNodeContainer((DataNodeContainer) childNode, moduleName, models, false, schemaContext);
             }
         }
 
     }
 
-
     /**
-     * Process the RPCs for a Module
-     * Spits out a file each of the name <rpcName>-input.json
-     * and <rpcName>-output.json for each RPC that contains
-     * input & output elements
+     * Process the RPCs for a Module Spits out a file each of the name <rpcName>-input.json and <rpcName>-output.json
+     * for each RPC that contains input & output elements
      *
      * @param module
      * @throws JSONException
      * @throws IOException
      */
-    private void processRPCs(final Module module, final JSONObject models) throws JSONException, IOException {
+    private void processRPCs(Module module, JSONObject models, SchemaContext schemaContext) throws JSONException,
+            IOException {
 
-        Set<RpcDefinition> rpcs =  module.getRpcs();
+        Set<RpcDefinition> rpcs = module.getRpcs();
         String moduleName = module.getName();
-        for(RpcDefinition rpc: rpcs) {
+        for (RpcDefinition rpc : rpcs) {
 
             ContainerSchemaNode input = rpc.getInput();
-            if(input!=null) {
-                JSONObject inputJSON = processContainer(input, moduleName, true, models);
-                String filename = rpc.getQName().getLocalName() + "-input";
+            if (input != null) {
+                JSONObject inputJSON = processDataNodeContainer(input, moduleName, models, schemaContext);
+                String filename = "(" + rpc.getQName().getLocalName() + ")input";
                 inputJSON.put("id", filename);
-                //writeToFile(filename, inputJSON.toString(2), moduleName);
+                // writeToFile(filename, inputJSON.toString(2), moduleName);
                 models.put(filename, inputJSON);
             }
 
             ContainerSchemaNode output = rpc.getOutput();
-            if(output!=null) {
-                JSONObject outputJSON = processContainer(output, moduleName, true, models);
-                String filename = rpc.getQName().getLocalName() + "-output";
+            if (output != null) {
+                JSONObject outputJSON = processDataNodeContainer(output, moduleName, models, schemaContext);
+                String filename = "(" + rpc.getQName().getLocalName() + ")output";
                 outputJSON.put("id", filename);
                 models.put(filename, outputJSON);
             }
         }
     }
 
+    /**
+     * Processes the 'identity' statement in a yang model and maps it to a 'model' in the Swagger JSON spec.
+     *
+     * @param module
+     *            The module from which the identity stmt will be processed
+     * @param models
+     *            The JSONObject in which the parsed identity will be put as a 'model' obj
+     * @throws JSONException
+     */
+    private void processIdentities(Module module, JSONObject models) throws JSONException {
+
+        String moduleName = module.getName();
+        Set<IdentitySchemaNode> idNodes = module.getIdentities();
+        LOG.debug("Processing Identities for module {} . Found {} identity statements", moduleName, idNodes.size());
+
+        for (IdentitySchemaNode idNode : idNodes) {
+            JSONObject identityObj = new JSONObject();
+            String identityName = idNode.getQName().getLocalName();
+            LOG.debug("Processing Identity: {}", identityName);
+
+            identityObj.put(ID_KEY, identityName);
+            identityObj.put(DESCRIPTION_KEY, idNode.getDescription());
+
+            JSONObject props = new JSONObject();
+            IdentitySchemaNode baseId = idNode.getBaseIdentity();
+
+            if (baseId == null) {
+                /**
+                 * This is a base identity. So lets see if it has sub types. If it does, then add them to the model
+                 * definition.
+                 */
+                Set<IdentitySchemaNode> derivedIds = idNode.getDerivedIdentities();
+
+                if (derivedIds != null) {
+                    JSONArray subTypes = new JSONArray();
+                    for (IdentitySchemaNode derivedId : derivedIds) {
+                        subTypes.put(derivedId.getQName().getLocalName());
+                    }
+                    identityObj.put(SUB_TYPES_KEY, subTypes);
+                }
+            } else {
+                /**
+                 * This is a derived entity. Add it's base type & move on.
+                 */
+                props.put(TYPE_KEY, baseId.getQName().getLocalName());
+            }
+
+            // Add the properties. For a base type, this will be an empty object as required by the Swagger spec.
+            identityObj.put(PROPERTIES_KEY, props);
+            models.put(identityName, identityObj);
+        }
+    }
 
     /**
-     * Processes the container node and populates the moduleJSON
+     * Processes the container and list nodes and populates the moduleJSON
      *
      * @param container
      * @param moduleName
+     * @param isConfig
      * @throws JSONException
      * @throws IOException
      */
-    private JSONObject processContainer(final ContainerSchemaNode container, final String moduleName, final boolean addSchemaStmt, final JSONObject models) throws JSONException, IOException{
-        JSONObject moduleJSON = getSchemaTemplate();
-        if(addSchemaStmt) {
-            moduleJSON = getSchemaTemplate();
-        } else {
-            moduleJSON = new JSONObject();
+    private JSONObject processDataNodeContainer(DataNodeContainer dataNode, String moduleName, JSONObject models,
+            SchemaContext schemaContext) throws JSONException, IOException {
+        return processDataNodeContainer(dataNode, moduleName, models, (Boolean) null, schemaContext);
+    }
+
+    private JSONObject processDataNodeContainer(DataNodeContainer dataNode, String moduleName, JSONObject models,
+            Boolean isConfig, SchemaContext schemaContext) throws JSONException, IOException {
+        if (dataNode instanceof ListSchemaNode || dataNode instanceof ContainerSchemaNode) {
+            Preconditions.checkArgument(dataNode instanceof SchemaNode, "Data node should be also schema node");
+            Iterable<DataSchemaNode> containerChildren = dataNode.getChildNodes();
+            JSONObject properties = processChildren(containerChildren, ((SchemaNode) dataNode).getQName(), moduleName,
+                    models, isConfig, schemaContext);
+
+            String nodeName = (BooleanUtils.isNotFalse(isConfig) ? OperationBuilder.CONFIG
+                    : OperationBuilder.OPERATIONAL) + ((SchemaNode) dataNode).getQName().getLocalName();
+
+            JSONObject childSchema = getSchemaTemplate();
+            childSchema.put(TYPE_KEY, OBJECT_TYPE);
+            childSchema.put(PROPERTIES_KEY, properties);
+            childSchema.put("id", nodeName);
+            models.put(nodeName, childSchema);
+
+            if (BooleanUtils.isNotFalse(isConfig)) {
+                createConcreteModelForPost(models, ((SchemaNode) dataNode).getQName().getLocalName(),
+                        createPropertiesForPost(dataNode));
+            }
+
+            JSONObject items = new JSONObject();
+            items.put(REF_KEY, nodeName);
+            JSONObject dataNodeProperties = new JSONObject();
+            dataNodeProperties.put(TYPE_KEY, dataNode instanceof ListSchemaNode ? ARRAY_TYPE : OBJECT_TYPE);
+            dataNodeProperties.put(ITEMS_KEY, items);
+
+            return dataNodeProperties;
         }
-        moduleJSON.put(TYPE_KEY, OBJECT_TYPE);
+        return null;
+    }
+
+    private void createConcreteModelForPost(final JSONObject models, final String localName, final JSONObject properties)
+            throws JSONException {
+        String nodePostName = OperationBuilder.CONFIG + localName + METHOD_NAME;
+        JSONObject postSchema = getSchemaTemplate();
+        postSchema.put(TYPE_KEY, OBJECT_TYPE);
+        postSchema.put("id", nodePostName);
+        postSchema.put(PROPERTIES_KEY, properties);
+        models.put(nodePostName, postSchema);
+    }
 
-        String containerDescription = container.getDescription();
-        moduleJSON.put(DESCRIPTION_KEY, containerDescription);
+    private JSONObject createPropertiesForPost(final DataNodeContainer dataNodeContainer) throws JSONException {
+        JSONObject properties = new JSONObject();
+        for (DataSchemaNode childNode : dataNodeContainer.getChildNodes()) {
+            if (childNode instanceof ListSchemaNode || childNode instanceof ContainerSchemaNode) {
+                JSONObject items = new JSONObject();
+                items.put(REF_KEY, "(config)" + childNode.getQName().getLocalName());
+                JSONObject property = new JSONObject();
+                property.put(TYPE_KEY, childNode instanceof ListSchemaNode ? ARRAY_TYPE : OBJECT_TYPE);
+                property.put(ITEMS_KEY, items);
+                properties.put(childNode.getQName().getLocalName(), property);
+            } else if (childNode instanceof LeafSchemaNode){
+                JSONObject property = processLeafNode((LeafSchemaNode)childNode);
+                properties.put(childNode.getQName().getLocalName(), property);
+            }
+        }
+        return properties;
+    }
 
-        Set<DataSchemaNode> containerChildren = container.getChildNodes();
-        JSONObject properties = processChildren(containerChildren, moduleName, models);
-        moduleJSON.put(PROPERTIES_KEY, properties);
-        return moduleJSON;
+    private JSONObject processChildren(Iterable<DataSchemaNode> nodes, QName parentQName, String moduleName,
+            JSONObject models, SchemaContext schemaContext) throws JSONException, IOException {
+        return processChildren(nodes, parentQName, moduleName, models, null, schemaContext);
     }
 
     /**
      * Processes the nodes
+     *
      * @param nodes
+     * @param parentQName
      * @param moduleName
+     * @param isConfig
      * @return
      * @throws JSONException
      * @throws IOException
      */
-    private JSONObject processChildren(final Set<DataSchemaNode> nodes, final String moduleName, final JSONObject models) throws JSONException, IOException {
+    private JSONObject processChildren(Iterable<DataSchemaNode> nodes, QName parentQName, String moduleName,
+            JSONObject models, Boolean isConfig, SchemaContext schemaContext) throws JSONException, IOException {
 
         JSONObject properties = new JSONObject();
 
-        for(DataSchemaNode node : nodes){
-            String name = node.getQName().getLocalName();
-            JSONObject property = null;
-            if(node instanceof LeafSchemaNode) {
-                property = processLeafNode((LeafSchemaNode)node);
-            } else if (node instanceof ListSchemaNode) {
-                property = processListSchemaNode((ListSchemaNode)node, moduleName, models);
+        for (DataSchemaNode node : nodes) {
+            if (isConfig == null || node.isConfiguration() == isConfig) {
 
-            } else if (node instanceof LeafListSchemaNode) {
-                property = processLeafListNode((LeafListSchemaNode)node);
+                String name = resolveNodesName(node, topLevelModule, schemaContext);
+                JSONObject property = null;
+                if (node instanceof LeafSchemaNode) {
+                    property = processLeafNode((LeafSchemaNode) node);
+                } else if (node instanceof ListSchemaNode) {
+                    property = processDataNodeContainer((ListSchemaNode) node, moduleName, models, isConfig,
+                            schemaContext);
 
-            } else if (node instanceof ChoiceNode) {
-                property = processChoiceNode((ChoiceNode)node, moduleName, models);
+                } else if (node instanceof LeafListSchemaNode) {
+                    property = processLeafListNode((LeafListSchemaNode) node);
 
-            } else if (node instanceof AnyXmlSchemaNode) {
-                property = processAnyXMLNode((AnyXmlSchemaNode)node);
+                } else if (node instanceof ChoiceSchemaNode) {
+                    property = processChoiceNode((ChoiceSchemaNode) node, moduleName, models, schemaContext);
 
-            } else if (node instanceof ContainerSchemaNode) {
-                property = processContainer((ContainerSchemaNode)node, moduleName, false, models);
+                } else if (node instanceof AnyXmlSchemaNode) {
+                    property = processAnyXMLNode((AnyXmlSchemaNode) node);
 
-            } else {
-                throw new IllegalArgumentException("Unknown DataSchemaNode type: " + node.getClass());
-            }
+                } else if (node instanceof ContainerSchemaNode) {
+                    property = processDataNodeContainer((ContainerSchemaNode) node, moduleName, models, isConfig,
+                            schemaContext);
 
-            property.putOpt(DESCRIPTION_KEY, node.getDescription());
-            properties.put(name, property);
+                } else {
+                    throw new IllegalArgumentException("Unknown DataSchemaNode type: " + node.getClass());
+                }
+
+                property.putOpt(DESCRIPTION_KEY, node.getDescription());
+                properties.put(name, property);
+            }
         }
         return properties;
     }
@@ -253,7 +377,7 @@ public class ModelGenerator {
      * @param listNode
      * @throws JSONException
      */
-    private JSONObject processLeafListNode(final LeafListSchemaNode listNode) throws JSONException {
+    private JSONObject processLeafListNode(LeafListSchemaNode listNode) throws JSONException {
         JSONObject props = new JSONObject();
         props.put(TYPE_KEY, ARRAY_TYPE);
 
@@ -274,14 +398,16 @@ public class ModelGenerator {
      * @throws JSONException
      * @throws IOException
      */
-    private JSONObject processChoiceNode(final ChoiceNode choiceNode, final String moduleName, final JSONObject models) throws JSONException, IOException {
+    private JSONObject processChoiceNode(ChoiceSchemaNode choiceNode, String moduleName, JSONObject models,
+            SchemaContext schemaContext) throws JSONException, IOException {
 
         Set<ChoiceCaseNode> cases = choiceNode.getCases();
 
         JSONArray choiceProps = new JSONArray();
-        for(ChoiceCaseNode choiceCase: cases) {
+        for (ChoiceCaseNode choiceCase : cases) {
             String choiceName = choiceCase.getQName().getLocalName();
-            JSONObject choiceProp = processChildren(choiceCase.getChildNodes(), moduleName, models);
+            JSONObject choiceProp = processChildren(choiceCase.getChildNodes(), choiceCase.getQName(), moduleName,
+                    models, schemaContext);
             JSONObject choiceObj = new JSONObject();
             choiceObj.put(choiceName, choiceProp);
             choiceObj.put(TYPE_KEY, OBJECT_TYPE);
@@ -295,78 +421,33 @@ public class ModelGenerator {
         return oneOfProps;
     }
 
-
     /**
      *
      * @param constraints
      * @param props
      * @throws JSONException
      */
-    private void processConstraints(final ConstraintDefinition constraints, final JSONObject props) throws JSONException {
+    private void processConstraints(ConstraintDefinition constraints, JSONObject props) throws JSONException {
         boolean isMandatory = constraints.isMandatory();
         props.put(REQUIRED_KEY, isMandatory);
 
         Integer minElements = constraints.getMinElements();
         Integer maxElements = constraints.getMaxElements();
-        if(minElements !=null) {
+        if (minElements != null) {
             props.put(MIN_ITEMS, minElements);
         }
-        if(maxElements !=null) {
+        if (maxElements != null) {
             props.put(MAX_ITEMS, maxElements);
         }
     }
 
-    /**
-     * Parses a ListSchema node.
-     *
-     * Due to a limitation of the RAML--->JAX-RS tool, sub-properties
-     * must be in a separate JSON schema file. Hence, we have to write
-     * some properties to a new file, while continuing to process the rest.
-     *
-     * @param listNode
-     * @param moduleName
-     * @return
-     * @throws JSONException
-     * @throws IOException
-     */
-    private JSONObject processListSchemaNode(final ListSchemaNode listNode, final String moduleName, final JSONObject models) throws JSONException, IOException {
-
-        Set<DataSchemaNode> listChildren = listNode.getChildNodes();
-        String fileName = listNode.getQName().getLocalName();
-
-        JSONObject childSchemaProperties = processChildren(listChildren, moduleName, models);
-        JSONObject childSchema = getSchemaTemplate();
-        childSchema.put(TYPE_KEY, OBJECT_TYPE);
-        childSchema.put(PROPERTIES_KEY, childSchemaProperties);
-
-        /*
-         * Due to a limitation of the RAML--->JAX-RS tool, sub-properties
-         * must be in a separate JSON schema file. Hence, we have to write
-         * some properties to a new file, while continuing to process the rest.
-         */
-        //writeToFile(fileName, childSchema.toString(2), moduleName);
-        childSchema.put("id", fileName);
-        models.put(fileName, childSchema);
-
-
-        JSONObject listNodeProperties = new JSONObject();
-        listNodeProperties.put(TYPE_KEY, ARRAY_TYPE);
-
-        JSONObject items = new JSONObject();
-        items.put(REF_KEY,fileName );
-        listNodeProperties.put(ITEMS_KEY, items);
-
-        return listNodeProperties;
-
-    }
-
     /**
      *
      * @param leafNode
      * @return
      * @throws JSONException
      */
-    private JSONObject processLeafNode(final LeafSchemaNode leafNode) throws JSONException {
+    private JSONObject processLeafNode(LeafSchemaNode leafNode) throws JSONException {
         JSONObject property = new JSONObject();
 
         String leafDescription = leafNode.getDescription();
@@ -384,7 +465,7 @@ public class ModelGenerator {
      * @return
      * @throws JSONException
      */
-    private JSONObject processAnyXMLNode(final AnyXmlSchemaNode leafNode) throws JSONException {
+    private JSONObject processAnyXMLNode(AnyXmlSchemaNode leafNode) throws JSONException {
         JSONObject property = new JSONObject();
 
         String leafDescription = leafNode.getDescription();
@@ -399,27 +480,27 @@ public class ModelGenerator {
      * @param property
      * @throws JSONException
      */
-    private void processTypeDef(final TypeDefinition<?> leafTypeDef, final JSONObject property) throws JSONException {
+    private void processTypeDef(TypeDefinition<?> leafTypeDef, JSONObject property) throws JSONException {
 
-        if(leafTypeDef instanceof ExtendedType){
+        if (leafTypeDef instanceof ExtendedType) {
             processExtendedType(leafTypeDef, property);
         } else if (leafTypeDef instanceof EnumerationType) {
-            processEnumType((EnumerationType)leafTypeDef, property);
+            processEnumType((EnumerationType) leafTypeDef, property);
 
         } else if (leafTypeDef instanceof BitsTypeDefinition) {
-            processBitsType((BitsTypeDefinition)leafTypeDef, property);
+            processBitsType((BitsTypeDefinition) leafTypeDef, property);
 
         } else if (leafTypeDef instanceof UnionTypeDefinition) {
-            processUnionType((UnionTypeDefinition)leafTypeDef, property);
+            processUnionType((UnionTypeDefinition) leafTypeDef, property);
 
         } else if (leafTypeDef instanceof IdentityrefTypeDefinition) {
-            property.putOpt(TYPE_KEY, "object");
+            property.putOpt(TYPE_KEY, ((IdentityrefTypeDefinition) leafTypeDef).getIdentity().getQName().getLocalName());
         } else if (leafTypeDef instanceof BinaryTypeDefinition) {
-            processBinaryType((BinaryTypeDefinition)leafTypeDef, property);
+            processBinaryType((BinaryTypeDefinition) leafTypeDef, property);
         } else {
-            //System.out.println("In else: " + leafTypeDef.getClass());
+            // System.out.println("In else: " + leafTypeDef.getClass());
             String jsonType = YANG_TYPE_TO_JSON_TYPE_MAPPING.get(leafTypeDef.getClass());
-            if(jsonType==null) {
+            if (jsonType == null) {
                 jsonType = "object";
             }
             property.putOpt(TYPE_KEY, jsonType);
@@ -432,29 +513,29 @@ public class ModelGenerator {
      * @param property
      * @throws JSONException
      */
-    private void processExtendedType(final TypeDefinition<?> leafTypeDef, final JSONObject property) throws JSONException {
+    private void processExtendedType(TypeDefinition<?> leafTypeDef, JSONObject property) throws JSONException {
         Object leafBaseType = leafTypeDef.getBaseType();
-        if(leafBaseType instanceof ExtendedType){
-            //recursively process an extended type until we hit a base type
-            processExtendedType((TypeDefinition<?>)leafBaseType, property);
+        if (leafBaseType instanceof ExtendedType) {
+            // recursively process an extended type until we hit a base type
+            processExtendedType((TypeDefinition<?>) leafBaseType, property);
         } else {
             List<LengthConstraint> lengthConstraints = ((ExtendedType) leafTypeDef).getLengthConstraints();
-            for(LengthConstraint lengthConstraint: lengthConstraints) {
+            for (LengthConstraint lengthConstraint : lengthConstraints) {
                 Number min = lengthConstraint.getMin();
                 Number max = lengthConstraint.getMax();
                 property.putOpt(MIN_LENGTH_KEY, min);
                 property.putOpt(MAX_LENGTH_KEY, max);
             }
             String jsonType = YANG_TYPE_TO_JSON_TYPE_MAPPING.get(leafBaseType.getClass());
-            property.putOpt(TYPE_KEY,jsonType );
+            property.putOpt(TYPE_KEY, jsonType);
         }
 
     }
 
     /*
-     *
-     */
-    private void processBinaryType(final BinaryTypeDefinition binaryType, final JSONObject property) throws JSONException {
+   *
+   */
+    private void processBinaryType(BinaryTypeDefinition binaryType, JSONObject property) throws JSONException {
         property.put(TYPE_KEY, STRING);
         JSONObject media = new JSONObject();
         media.put(BINARY_ENCODING_KEY, BASE_64);
@@ -467,10 +548,10 @@ public class ModelGenerator {
      * @param property
      * @throws JSONException
      */
-    private void processEnumType(final EnumerationType enumLeafType, final JSONObject property) throws JSONException {
+    private void processEnumType(EnumerationType enumLeafType, JSONObject property) throws JSONException {
         List<EnumPair> enumPairs = enumLeafType.getValues();
         List<String> enumNames = new ArrayList<String>();
-        for(EnumPair enumPair: enumPairs) {
+        for (EnumPair enumPair : enumPairs) {
             enumNames.add(enumPair.getName());
         }
         property.putOpt(ENUM, new JSONArray(enumNames));
@@ -482,14 +563,14 @@ public class ModelGenerator {
      * @param property
      * @throws JSONException
      */
-    private void processBitsType(final BitsTypeDefinition bitsType, final JSONObject property) throws JSONException{
+    private void processBitsType(BitsTypeDefinition bitsType, JSONObject property) throws JSONException {
         property.put(TYPE_KEY, ARRAY_TYPE);
         property.put(MIN_ITEMS, 0);
         property.put(UNIQUE_ITEMS_KEY, true);
         JSONArray enumValues = new JSONArray();
 
         List<Bit> bits = bitsType.getBits();
-        for(Bit bit: bits) {
+        for (Bit bit : bits) {
             enumValues.put(bit.getName());
         }
         JSONObject itemsValue = new JSONObject();
@@ -497,27 +578,28 @@ public class ModelGenerator {
         property.put(ITEMS_KEY, itemsValue);
     }
 
-
     /**
      *
      * @param unionType
      * @param property
      * @throws JSONException
      */
-    private void processUnionType(final UnionTypeDefinition unionType, final JSONObject property) throws JSONException{
+    private void processUnionType(UnionTypeDefinition unionType, JSONObject property) throws JSONException {
 
-        List<TypeDefinition<?>> unionTypes = unionType.getTypes();
-        JSONArray unionArray = new JSONArray();
-        for(TypeDefinition<?> typeDef: unionTypes) {
-            unionArray.put(YANG_TYPE_TO_JSON_TYPE_MAPPING.get(typeDef.getClass()));
+        StringBuilder type = new StringBuilder();
+        for (TypeDefinition<?> typeDef : unionType.getTypes()) {
+            if (type.length() > 0) {
+                type.append(" or ");
+            }
+            type.append(YANG_TYPE_TO_JSON_TYPE_MAPPING.get(typeDef.getClass()));
         }
-        property.put(TYPE_KEY, unionArray);
-    }
 
+        property.put(TYPE_KEY, type);
+    }
 
     /**
-     * Helper method to generate a pre-filled
-     * JSON schema object.
+     * Helper method to generate a pre-filled JSON schema object.
+     *
      * @return
      * @throws JSONException
      */
@@ -527,4 +609,5 @@ public class ModelGenerator {
 
         return schemaJSON;
     }
+
 }