diff --git a/native/src/main/java/io/ballerina/stdlib/mi/OMElementConverter.java b/native/src/main/java/io/ballerina/stdlib/mi/OMElementConverter.java
index 4cfb86f..dfa70b2 100644
--- a/native/src/main/java/io/ballerina/stdlib/mi/OMElementConverter.java
+++ b/native/src/main/java/io/ballerina/stdlib/mi/OMElementConverter.java
@@ -76,13 +76,7 @@ private static QName getQNameOMElement(OMElement omElement) {
     }
 
     private static QName getQNameOMAttribute(OMAttribute omAttribute) {
-        // If prefix is not provided test case will fail
-        if (omAttribute.getPrefix() == null) {
-            return new QName(omAttribute.getNamespaceURI(), omAttribute.getLocalName());
-        } else {
-            return new QName(omAttribute.getNamespaceURI(), omAttribute.getLocalName(), omAttribute.getPrefix());
-        }
-
+        return omAttribute.getQName();
     }
 
     private static void addAttributes(OMElement omElement, BXmlItem xmlItem) {
@@ -99,10 +93,8 @@ private static void addAttributes(OMElement omElement, BXmlItem xmlItem) {
             //CHECK: Good to put break point here and check the values
             attributesMap.put(StringUtils.fromString(qName.toString()),
                     StringUtils.fromString(attribute.getAttributeValue()));
-            if (attribute.getPrefix() != null) {
-                if (!attribute.getPrefix().isEmpty()) {
-                    usedNS.add(qName);
-                }
+            if (!qName.getPrefix().isEmpty()) {
+                usedNS.add(qName);
             }
         }
         if (omElement.getQName().getPrefix() != null) {
diff --git a/tests/src/test/java/org/ballerina/test/TestTransform.java b/tests/src/test/java/org/ballerina/test/TestTransform.java
index c987101..d7a1f23 100644
--- a/tests/src/test/java/org/ballerina/test/TestTransform.java
+++ b/tests/src/test/java/org/ballerina/test/TestTransform.java
@@ -294,6 +294,78 @@ public void testJsonTransform1() {
                 "{\"name\":\"John\",\"age\":23,\"city\":\"Colombo\",\"country\":\"Sri Lanka\"}");
     }
 
+    @Test
+    public void testXmlReturn1() throws XMLStreamException {
+        String project = "project3";
+        ModuleInfo moduleInfo = new ModuleInfo("testOrg", project, "1");
+        Mediator mediator = new Mediator(moduleInfo);
+
+        List<Property> properties = new ArrayList<>() {{
+            add(new Property("paramFunctionName", "testXmlReturn"));
+            add(new Property("paramSize", 1));
+            add(new Property("param0", "x"));
+            add(new Property("paramType0", "xml"));
+            add(new Property("returnType", "xml"));
+        }};
+        TestMessageContext context = createMessageContext(properties);
+
+        Stack<TemplateContext> stack = new Stack<>();
+        TemplateContext templateContext = new TemplateContext("testTemplateFunc", new ArrayList<>());
+        stack.push(templateContext);
+
+        HashMap<Object, Object> map = new HashMap<>();
+        map.put("x", AXIOMUtil.stringToOM("<Document><FIToFICstmrCdtTrf><GrpHdr><MsgId>123489</MsgId><CreDtTm" +
+                ">0111522180</CreDtTm><NbOfTxs>1</NbOfTxs><SttlmInf><SttlmMtd>INDA</SttlmMtd></SttlmInf></GrpHdr" +
+                "><CdtTrfTxInf><PmtId><EndToEndId>12348912a456789123</EndToEndId></PmtId><IntrBkSttlmAmt " +
+                "Ccy=\"USD\">500000</IntrBkSttlmAmt><ChrgBr>DEBT</ChrgBr><Dbtr><number>1</number></Dbtr><CdtrAgt" +
+                "><FinInstnId>100009</FinInstnId></CdtrAgt></CdtTrfTxInf></FIToFICstmrCdtTrf></Document>"));
+        map.put("Result", "r");
+        templateContext.setMappedValues(map);
+
+        context.setProperty("_SYNAPSE_FUNCTION_STACK", stack);
+        mediator.mediate(context);
+        Assert.assertEquals(context.getProperty("r").toString(),
+                "<Document><FIToFICstmrCdtTrf><GrpHdr><MsgId>123489</MsgId><CreDtTm>0111522180</CreDtTm><NbOfTxs>1" +
+                        "</NbOfTxs><SttlmInf><SttlmMtd>INDA</SttlmMtd></SttlmInf></GrpHdr><CdtTrfTxInf><PmtId" +
+                        "><EndToEndId>12348912a456789123</EndToEndId></PmtId><IntrBkSttlmAmt " +
+                        "Ccy=\"USD\">500000</IntrBkSttlmAmt><ChrgBr>DEBT</ChrgBr><Dbtr><number>1</number></Dbtr" +
+                        "><CdtrAgt><FinInstnId>100009</FinInstnId></CdtrAgt></CdtTrfTxInf></FIToFICstmrCdtTrf" +
+                        "></Document>");
+    }
+
+    @Test
+    public void testXmlReturn2() throws XMLStreamException {
+        String project = "project3";
+        ModuleInfo moduleInfo = new ModuleInfo("testOrg", project, "1");
+        Mediator mediator = new Mediator(moduleInfo);
+
+        List<Property> properties = new ArrayList<>() {{
+            add(new Property("paramFunctionName", "testXmlReturn"));
+            add(new Property("paramSize", 1));
+            add(new Property("param0", "x"));
+            add(new Property("paramType0", "xml"));
+            add(new Property("returnType", "xml"));
+        }};
+        TestMessageContext context = createMessageContext(properties);
+
+        Stack<TemplateContext> stack = new Stack<>();
+        TemplateContext templateContext = new TemplateContext("testTemplateFunc", new ArrayList<>());
+        stack.push(templateContext);
+
+        HashMap<Object, Object> map = new HashMap<>();
+        map.put("x", AXIOMUtil.stringToOM("<library xmlns:bk=\"http://example.com/book\" xmlns:auth=\"http://example" +
+                ".com/author\"><bk:Book bk:id=\"001\"><bk:Title>1984</bk:Title><bk:Author " +
+                "auth:country=\"UK\"><auth:Name>George Orwell</auth:Name></bk:Author></bk:Book></library>"));
+        map.put("Result", "r");
+        templateContext.setMappedValues(map);
+
+        context.setProperty("_SYNAPSE_FUNCTION_STACK", stack);
+        mediator.mediate(context);
+        Assert.assertEquals(context.getProperty("r").toString(),
+                "<library><bk:Book xmlns:bk=\"http://example.com/book\" " +
+                        "bk:id=\"001\"><bk:Title>1984</bk:Title><bk:Author xmlns:auth=\"http://example.com/author\" " +
+                        "auth:country=\"UK\"><auth:Name>George Orwell</auth:Name></bk:Author></bk:Book></library>");
+    }
 
     private TestMessageContext createMessageContext(List<Property> properties) {
         TestMessageContext context = new TestMessageContext();
diff --git a/tests/src/test/resources/ballerina/project3/main.bal b/tests/src/test/resources/ballerina/project3/main.bal
index 956bd4f..cee5007 100644
--- a/tests/src/test/resources/ballerina/project3/main.bal
+++ b/tests/src/test/resources/ballerina/project3/main.bal
@@ -1,3 +1,7 @@
 public function test(xml xmlA, xml xmlB, xml xmlC, xml xmlD) returns xml {
     return xml `<may22>${xmlA}${xmlB}${xmlC}${xmlD}</may22>`;
 }
+
+public function testXmlReturn(xml x) returns xml {
+    return x;
+}