Bug 1754161 - [marionette] Merge Marionette::_unwrap_response in _from_json r=webdriver-reviewers,jgraham

Depends on D138159

_from_json should be redundant with _unwrap_response now.

Differential Revision: https://phabricator.services.mozilla.com/D138160

testing/marionette/client/marionette_driver/marionette.py

@@ -378,7 +378,7 @@ class HTMLElement(object):
 class ShadowRoot(object):
     """A Class to handling Shadow Roots"""
 
-    identifiers = WEB_SHADOW_ROOT_KEY
+    identifiers = (WEB_SHADOW_ROOT_KEY,)
 
     def __init__(self, marionette, id, kind=WEB_SHADOW_ROOT_KEY):
         self.marionette = marionette
@@ -668,24 +668,9 @@ class Marionette(object):
             self._handle_error(err)
 
         if key is not None:
-            return self._unwrap_response(res.get(key))
+            return self._from_json(res.get(key))
         else:
-            return self._unwrap_response(res)
-
-    def _unwrap_response(self, value):
-
-        if isinstance(value, dict) and any(
-            k in value.keys() for k in HTMLElement.identifiers
-        ):
-            return HTMLElement._from_json(value, self)
-        elif isinstance(value, dict) and ShadowRoot.identifiers in value.keys():
-            return ShadowRoot._from_json(value, self)
-        elif isinstance(value, dict):
-            return {key: self._unwrap_response(val) for key, val in value.items()}
-        elif isinstance(value, list):
-            return list(self._unwrap_response(item) for item in value)
-        else:
-            return value
+            return self._from_json(res)
 
     def _handle_error(self, obj):
         error = obj["error"]
@@ -1599,19 +1584,18 @@ class Marionette(object):
         return wrapped
 
     def _from_json(self, value):
-        if isinstance(value, list):
-            unwrapped = []
-            for item in value:
-                unwrapped.append(self._from_json(item))
-            return unwrapped
+        if isinstance(value, dict) and any(
+            k in value.keys() for k in HTMLElement.identifiers
+        ):
+            return HTMLElement._from_json(value, self)
+        elif isinstance(value, dict) and any(
+            k in value.keys() for k in ShadowRoot.identifiers
+        ):
+            return ShadowRoot._from_json(value, self)
         elif isinstance(value, dict):
-            unwrapped = {}
-            for key in value:
-                if key in HTMLElement.identifiers:
-                    return HTMLElement._from_json(value[key], self)
-                else:
-                    unwrapped[key] = self._from_json(value[key])
-            return unwrapped
+            return {key: self._from_json(val) for key, val in value.items()}
+        elif isinstance(value, list):
+            return list(self._from_json(item) for item in value)
         else:
             return value
 
@@ -1712,7 +1696,7 @@ class Marionette(object):
             if script_timeout is not None:
                 self.timeout.script = original_timeout
 
-        return self._from_json(rv)
+        return rv
 
     def execute_async_script(
         self,
@@ -1782,7 +1766,7 @@ class Marionette(object):
             if script_timeout is not None:
                 self.timeout.script = original_timeout
 
-        return self._from_json(rv)
+        return rv
 
     def find_element(self, method, target, id=None):
         """Returns an :class:`~marionette_driver.marionette.HTMLElement`

testing/marionette/harness/marionette_harness/tests/unit/test_execute_script.py

@@ -210,6 +210,26 @@ class TestExecuteContent(MarionetteTestCase):
         )
         self.assertEqual(expected, actual)
 
+    def test_return_web_element_nested_array(self):
+        self.marionette.navigate(elements)
+        expected = self.marionette.find_elements(By.TAG_NAME, "p")
+        actual = self.marionette.execute_script(
+            """
+            let els = document.querySelectorAll('p')
+            return { els: [els[0], els[1]] }"""
+        )
+        self.assertEqual(expected, actual["els"])
+
+    def test_return_web_element_nested_dict(self):
+        self.marionette.navigate(elements)
+        expected = self.marionette.find_element(By.TAG_NAME, "p")
+        actual = self.marionette.execute_script(
+            """
+            let el = document.querySelector('p')
+            return { path: { to: { el } } }"""
+        )
+        self.assertEqual(expected, actual["path"]["to"]["el"])
+
     # Bug 938228 identifies a problem with unmarshaling NodeList
     # objects from the DOM.  document.querySelectorAll returns this
     # construct.
@@ -482,6 +502,12 @@ class TestExecuteChrome(WindowManagerMixin, TestExecuteContent):
     def test_return_web_element_array(self):
         pass
 
+    def test_return_web_element_nested_array(self):
+        pass
+
+    def test_return_web_element_nested_dict(self):
+        pass
+
     def test_return_web_element_nodelist(self):
         pass