From ff2b265ba13716f01ab82e6fa5fccac893063c89 Mon Sep 17 00:00:00 2001
From: David Burns <dburns@mozilla.com>
Date: Fri, 11 Sep 2015 23:55:59 +0100
Subject: [PATCH] Bug 1137972: Centralize unwrapping the response from
 marionette; r=jgriffin

--HG--
extra : commitid : 94hIuseVTZ7
extra : rebase_source : 15d3704629e1caccd483ed075e7ede0ce63d6203
---
 .../driver/marionette_driver/marionette.py    | 30 ++++++++++++-------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/testing/marionette/driver/marionette_driver/marionette.py b/testing/marionette/driver/marionette_driver/marionette.py
index 37ffff203fae..61ca96d4692f 100644
--- a/testing/marionette/driver/marionette_driver/marionette.py
+++ b/testing/marionette/driver/marionette_driver/marionette.py
@@ -713,8 +713,22 @@ class Marionette(object):
             self._handle_error(resp)
 
         if key is not None:
-            resp = resp[key]
-        return resp
+            return self._unwrap_response(resp[key])
+        else:
+            return self._unwrap_response(resp)
+
+    def _unwrap_response(self, value):
+        resp = ""
+        if isinstance(value, dict) and \
+        ('ELEMENT' in value or 'element-6066-11e4-a52e-4f735466cecf' in value):
+            if value.get('ELEMENT'):
+                return HTMLElement(self, value.get('ELEMENT'))
+            else:
+                return HTMLElement(self, value.get('element-6066-11e4-a52e-4f735466cecf'))
+        elif isinstance(value, list):
+            return list(self._unwrap_response(item) for item in value)
+        else:
+            return value
 
     def _emulator_cmd(self, id, cmd):
         if not self.emulator:
@@ -1595,9 +1609,7 @@ class Marionette(object):
         body = {"value": target, "using": method}
         if id:
             body["element"] = id
-        el = self._send_message("findElement", body, key="value")
-        ref = el["ELEMENT"]
-        return HTMLElement(self, ref)
+        return self._send_message("findElement", body, key="value")
 
     def find_elements(self, method, target, id=None):
         """Returns a list of all HTMLElement instances that match the
@@ -1622,13 +1634,9 @@ class Marionette(object):
         body = {"value": target, "using": method}
         if id:
             body["element"] = id
-        els = self._send_message(
+        return self._send_message(
             "findElements", body, key="value" if self.protocol == 1 else None)
-        assert(isinstance(els, list))
-        rv = []
-        for el in els:
-            rv.append(HTMLElement(self, el["ELEMENT"]))
-        return rv
+
 
     def get_active_element(self):
         el = self._send_message("getActiveElement", key="value")