[ONOS-6951] Union with two or more identities code generation fix.

Change-Id: If61fa8e2a6ce1e963c4d4c089d4bd16fe7772ee9
diff --git a/compiler/base/translator/src/main/java/org/onosproject/yang/compiler/translator/tojava/TempJavaTypeFragmentFiles.java b/compiler/base/translator/src/main/java/org/onosproject/yang/compiler/translator/tojava/TempJavaTypeFragmentFiles.java
index 174c0d8..81a7718 100644
--- a/compiler/base/translator/src/main/java/org/onosproject/yang/compiler/translator/tojava/TempJavaTypeFragmentFiles.java
+++ b/compiler/base/translator/src/main/java/org/onosproject/yang/compiler/translator/tojava/TempJavaTypeFragmentFiles.java
@@ -16,6 +16,7 @@
 
 package org.onosproject.yang.compiler.translator.tojava;
 
+import org.onosproject.yang.compiler.datamodel.YangIdentityRef;
 import org.onosproject.yang.compiler.datamodel.YangNode;
 import org.onosproject.yang.compiler.datamodel.YangType;
 import org.onosproject.yang.compiler.datamodel.YangTypeDef;
@@ -285,6 +286,7 @@
         if (yangTypeHolder instanceof YangUnion) {
             verifyUnionTypes(typeList, yangTypeHolder);
         }
+        addIdentityToImport(typeList, config);
         if (typeList != null) {
             List<YangType<?>> types = validateTypes(typeList);
             for (YangType<?> type : types) {
@@ -326,6 +328,17 @@
         }
     }
 
+    //TODO: Union with two identities code generation has to be changed.
+    private void addIdentityToImport(List<YangType<?>> types,
+                                     YangPluginConfig config) {
+        for (YangType<?> type :types) {
+            Object ex = type.getDataTypeExtendedInfo();
+            if (ex != null && ex instanceof YangIdentityRef) {
+                getAttributeForType(type, config);
+            }
+        }
+    }
+
     /**
      * Returns java attribute.
      *
diff --git a/compiler/plugin/maven/src/test/java/org/onosproject/yang/compiler/plugin/maven/AugmentTranslatorTest.java b/compiler/plugin/maven/src/test/java/org/onosproject/yang/compiler/plugin/maven/AugmentTranslatorTest.java
index 66dc468..e233495 100644
--- a/compiler/plugin/maven/src/test/java/org/onosproject/yang/compiler/plugin/maven/AugmentTranslatorTest.java
+++ b/compiler/plugin/maven/src/test/java/org/onosproject/yang/compiler/plugin/maven/AugmentTranslatorTest.java
@@ -562,4 +562,35 @@
         YangPluginConfig.compileCode(COMP);
         deleteDirectory(DIR);
     }
+
+    /**
+     * Checks union with more that two identities for proper code generation.
+     *
+     * @throws IOException            if any error occurs during IO on files
+     * @throws ParserException        if any error occurs during parsing
+     * @throws MojoExecutionException if any mojo operation fail
+     */
+    @Test
+    public void processUnionIdentity() throws IOException,
+            ParserException, MojoExecutionException {
+
+        deleteDirectory(DIR);
+        String dir = "src/test/resources/unionTranslator/unionidentity";
+
+        Set<Path> paths = new HashSet<>();
+        for (String file : getYangFiles(dir)) {
+            paths.add(Paths.get(file));
+        }
+
+        utilManager.createYangFileInfoSet(paths);
+        utilManager.parseYangFileInfoSet();
+        utilManager.createYangNodeSet();
+        utilManager.resolveDependenciesUsingLinker();
+
+        YangPluginConfig yangPluginConfig = new YangPluginConfig();
+        yangPluginConfig.setCodeGenDir(DIR);
+        utilManager.translateToJava(yangPluginConfig);
+        YangPluginConfig.compileCode(COMP);
+        deleteDirectory(DIR);
+    }
 }
diff --git a/compiler/plugin/maven/src/test/resources/unionTranslator/unionidentity/module1.yang b/compiler/plugin/maven/src/test/resources/unionTranslator/unionidentity/module1.yang
new file mode 100644
index 0000000..1a58f44
--- /dev/null
+++ b/compiler/plugin/maven/src/test/resources/unionTranslator/unionidentity/module1.yang
@@ -0,0 +1,29 @@
+module module1 {
+    yang-version 1;
+    namespace "http://huawei.com:mod1";
+    prefix Ant;
+
+    import module2 {
+        prefix mod2;
+    }
+
+    typedef def6 {
+        type binary;
+    }
+
+    leaf icon {
+        type union {
+            type identityref {
+                base mod2:axis;
+            }
+            type identityref {
+                base mod2:represent;
+            }
+            type identityref {
+                base mod2:area;
+            }
+            type uint8;
+            type def6;
+        }
+    }
+}
\ No newline at end of file
diff --git a/compiler/plugin/maven/src/test/resources/unionTranslator/unionidentity/module2.yang b/compiler/plugin/maven/src/test/resources/unionTranslator/unionidentity/module2.yang
new file mode 100644
index 0000000..3328252
--- /dev/null
+++ b/compiler/plugin/maven/src/test/resources/unionTranslator/unionidentity/module2.yang
@@ -0,0 +1,39 @@
+module module2 {
+    yang-version 1;
+    namespace "http://huawei.com1:mod2";
+
+    prefix mod2;
+
+    identity axis {
+    }
+
+    identity left {
+        base axis;
+    }
+
+    identity right {
+        base axis;
+    }
+
+    identity represent {
+    }
+
+    identity horizontal {
+        base represent;
+    }
+
+    identity vertical {
+        base represent;
+    }
+
+    identity area {
+    }
+
+    identity square {
+        base area;
+    }
+
+    identity rectangle {
+        base area;
+    }
+}
\ No newline at end of file