diff --git a/flow-tests/test-root-context/src/test/java/com/vaadin/flow/uitest/ui/template/ChildOrderIT.java b/flow-tests/test-root-context/src/test/java/com/vaadin/flow/uitest/ui/template/ChildOrderIT.java
index 7433d01482e..eef2b309cc6 100644
--- a/flow-tests/test-root-context/src/test/java/com/vaadin/flow/uitest/ui/template/ChildOrderIT.java
+++ b/flow-tests/test-root-context/src/test/java/com/vaadin/flow/uitest/ui/template/ChildOrderIT.java
@@ -24,7 +24,10 @@
 import org.openqa.selenium.By;
 import org.openqa.selenium.WebElement;
 
+import com.vaadin.flow.component.html.testbench.DivElement;
+import com.vaadin.flow.component.html.testbench.NativeButtonElement;
 import com.vaadin.flow.testutil.ChromeBrowserTest;
+import com.vaadin.testbench.TestBenchElement;
 
 /**
  * Tests to validate the ordering of server-side nodes when added alongside
@@ -34,19 +37,19 @@
  */
 public class ChildOrderIT extends ChromeBrowserTest {
 
-    private WebElement root;
+    private TestBenchElement root;
 
     @Before
     public void init() {
         open();
-        waitForElementPresent(By.id("root"));
-        root = findElement(By.id("root"));
+        waitForElementPresent(By.tagName("child-order-template"));
+        root = $("child-order-template").first();
     }
 
     @Test
     public void appendElementsFromServer_elementsAreAddedAfterExistingOnes() {
-        WebElement container = findInShadowRoot(root,
-                By.id("containerWithElement")).get(0);
+        DivElement container = root.$(DivElement.class)
+                .id("containerWithElement");
 
         assertNodeOrder(container, "Client child");
 
@@ -83,8 +86,8 @@ public void appendElementsFromServer_elementsAreAddedAfterExistingOnes() {
 
     @Test
     public void prependElementsFromServer_elementsAreAddedBeforeExistingOnes() {
-        WebElement container = findInShadowRoot(root,
-                By.id("containerWithElement")).get(0);
+        DivElement container = root.$(DivElement.class)
+                .id("containerWithElement");
 
         assertNodeOrder(container, "Client child");
 
@@ -117,8 +120,7 @@ public void prependElementsFromServer_elementsAreAddedBeforeExistingOnes() {
 
     @Test
     public void appendTextsFromServer_textsAreAddedAfterExistingOnes() {
-        WebElement container = findInShadowRoot(root,
-                By.id("containerWithText")).get(0);
+        DivElement container = root.$(DivElement.class).id("containerWithText");
 
         assertNodeOrder(container, "Client text");
 
@@ -155,8 +157,7 @@ public void appendTextsFromServer_textsAreAddedAfterExistingOnes() {
 
     @Test
     public void prependTextsFromServer_textsAreAddedBeforeExistingOnes() {
-        WebElement container = findInShadowRoot(root,
-                By.id("containerWithText")).get(0);
+        DivElement container = root.$(DivElement.class).id("containerWithText");
 
         assertNodeOrder(container, "Client text");
 
@@ -189,8 +190,8 @@ public void prependTextsFromServer_textsAreAddedBeforeExistingOnes() {
 
     @Test
     public void containerWithElementAddedOnConstructor_orderIsPreserved() {
-        WebElement container = findInShadowRoot(root,
-                By.id("containerWithElementAddedOnConstructor")).get(0);
+        DivElement container = root.$(DivElement.class)
+                .id("containerWithElementAddedOnConstructor");
 
         assertNodeOrder(container, "Client child", "Server child 1",
                 "Server child 2");
@@ -199,7 +200,8 @@ public void containerWithElementAddedOnConstructor_orderIsPreserved() {
     private void clickAndWaitForContainerToChange(WebElement container,
             String buttonToclick) {
         String innertText = container.getAttribute("innerText");
-        WebElement button = findInShadowRoot(root, By.id(buttonToclick)).get(0);
+        NativeButtonElement button = root.$(NativeButtonElement.class)
+                .id(buttonToclick);
         button.click();
         waitUntilNot(driver -> container.getAttribute("innerText")
                 .equals(innertText));