From b2022dfd64e28f5863b7e5a211e54ef265cf46b2 Mon Sep 17 00:00:00 2001
From: Artur <artur@vaadin.com>
Date: Mon, 12 Feb 2018 08:48:00 +0200
Subject: [PATCH] Fix ElementQuery for built-in elements or elements without a
 shadow root (#990)

Fixes #989
---
 .../com/vaadin/testbench/ElementQuery.java     | 18 +++++++++---------
 .../java/com/vaadin/tests/ElementQueryIT.java  | 10 ++++++++++
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/vaadin-testbench-core/src/main/java/com/vaadin/testbench/ElementQuery.java b/vaadin-testbench-core/src/main/java/com/vaadin/testbench/ElementQuery.java
index 523ea7380..4afadd7a4 100644
--- a/vaadin-testbench-core/src/main/java/com/vaadin/testbench/ElementQuery.java
+++ b/vaadin-testbench-core/src/main/java/com/vaadin/testbench/ElementQuery.java
@@ -259,15 +259,15 @@ private List<T> executeSearch(Integer index) {
         TestBenchElement elementContext;
         JavascriptExecutor executor;
         if (getContext() instanceof TestBenchElement) {
-            script.append(
-                    "var result = arguments[0].shadowRoot.querySelectorAll(arguments[1]+arguments[2]);");
-            script.append(
-                    "var light = arguments[0].querySelectorAll(arguments[1]+arguments[2]);");
-            script.append("if (light.length > 0) {");
-            script.append(
-                    "result = Array.prototype.slice.call(result).concat(Array.prototype.slice.call(light));");
-            script.append("}");
-            script.append("return result");
+            script.append("var result = [];" //
+                    + "if (arguments[0].shadowRoot) {" //
+                    + "  var shadow = arguments[0].shadowRoot.querySelectorAll(arguments[1]+arguments[2]);" //
+                    + "  result = result.concat(Array.prototype.slice.call(shadow));" //
+                    + "}" //
+                    + "var light = arguments[0].querySelectorAll(arguments[1]+arguments[2]);" //
+                    + "result = result.concat(Array.prototype.slice.call(light));" //
+                    + "return result" //
+            );
             elementContext = (TestBenchElement) getContext();
             executor = elementContext.getCommandExecutor().getDriver();
         } else if (getContext() instanceof WebDriver) {
diff --git a/vaadin-testbench-integration-tests/src/test/java/com/vaadin/tests/ElementQueryIT.java b/vaadin-testbench-integration-tests/src/test/java/com/vaadin/tests/ElementQueryIT.java
index b8e558f68..7ea0ca35d 100644
--- a/vaadin-testbench-integration-tests/src/test/java/com/vaadin/tests/ElementQueryIT.java
+++ b/vaadin-testbench-integration-tests/src/test/java/com/vaadin/tests/ElementQueryIT.java
@@ -99,4 +99,14 @@ public void findTestBenchElement() throws Exception {
                 .waitForFirst().$(TestBenchElement.class).id("shadow-button-2");
         Assert.assertNotNull(button);
     }
+
+    @Test
+    public void findTestBenchElementChild() throws Exception {
+        openTestURL();
+
+        TestBenchElement button = $(PolymerTemplateViewElement.class)
+                .waitForFirst().$(TestBenchElement.class).first()
+                .$(TestBenchElement.class).first();
+        Assert.assertEquals("Shadow Button 1", button.getText());
+    }
 }