diff --git a/app/hydrator/services/create/stores/config-store.js b/app/hydrator/services/create/stores/config-store.js
index 61426afee3e..4bb078215d6 100644
--- a/app/hydrator/services/create/stores/config-store.js
+++ b/app/hydrator/services/create/stores/config-store.js
@@ -207,7 +207,7 @@ class HydratorPlusPlusConfigStore {
           _backendProperties: node._backendProperties,
         },
         information: node.information,
-        outputSchema: angular.isArray(node.outputSchema) ? node.outputSchema[0].schema : node.outputSchema,
+        outputSchema: node.outputSchema,
         inputSchema: node.inputSchema
       };
 
diff --git a/src/e2e-test/java/io/cdap/cdap/ui/stepsdesign/PipelineEdit.java b/src/e2e-test/java/io/cdap/cdap/ui/stepsdesign/PipelineEdit.java
index 380b716aa51..83db510e330 100644
--- a/src/e2e-test/java/io/cdap/cdap/ui/stepsdesign/PipelineEdit.java
+++ b/src/e2e-test/java/io/cdap/cdap/ui/stepsdesign/PipelineEdit.java
@@ -195,6 +195,8 @@ public void draftShouldNotExist(String pipelineName) {
 
   @Then("Generated runtime arguments should be empty")
   public void checkRuntimeArgsEmpty() {
+    WaitHelper.waitForElementToBePresent(
+      By.cssSelector(Helper.getCssSelectorByDataTestId("generated-runtimeargs-count")));
     Assert.assertTrue(Helper.locateElementByTestId("generated-runtimeargs-count").getText().charAt(0) == '0');
     ElementHelper.clickOnElement(Helper.locateElementByTestId("generated-runtimeargs"));
     List<WebElement> keys = Helper.locateElementsByTestId("runtimeargs-key");
@@ -209,6 +211,8 @@ public void checkRuntimeArgsEmpty() {
 
   @Then("Generated runtime arguments should not be empty")
   public void checkRuntimeArgsNotEmpty() {
+    WaitHelper.waitForElementToBePresent(
+      By.cssSelector(Helper.getCssSelectorByDataTestId("generated-runtimeargs-count")));
     Assert.assertFalse(Helper.locateElementByTestId("generated-runtimeargs-count").getText().charAt(0) == '0');
     ElementHelper.clickOnElement(Helper.locateElementByTestId("generated-runtimeargs"));
     List<WebElement> keys = Helper.locateElementsByTestId("runtimeargs-key");