We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Socket errors thrown from device obj are not marked with device tag
We should wrap `socket.error` with a child class of `DeviceError`.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mobly/controllers/android_device_lib/jsonrpc_client_base.py`
Content:
```
1 # Copyright 2016 Google Inc.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """Base class for clients that communicate with apps over a JSON RPC interface.
15 
16 The JSON protocol expected by this module is:
17 
18 .. code-block:: json
19 
20     Request:
21     {
22         "id": <monotonically increasing integer containing the ID of 
23                this request>
24         "method": <string containing the name of the method to execute>
25         "params": <JSON array containing the arguments to the method>
26     }
27 
28     Response:
29     {
30         "id": <int id of request that this response maps to>,
31         "result": <Arbitrary JSON object containing the result of
32                    executing the method. If the method could not be
33                    executed or returned void, contains 'null'.>,
34         "error": <String containing the error thrown by executing the
35                   method. If no error occurred, contains 'null'.>
36         "callback": <String that represents a callback ID used to
37                      identify events associated with a particular
38                      CallbackHandler object.>
39     }
40 """
41 
42 from builtins import str
43 
44 import json
45 import logging
46 import socket
47 import threading
48 
49 from mobly.controllers.android_device_lib import callback_handler
50 from mobly.controllers.android_device_lib import errors
51 
52 # UID of the 'unknown' jsonrpc session. Will cause creation of a new session.
53 UNKNOWN_UID = -1
54 
55 # Maximum time to wait for the socket to open on the device.
56 _SOCKET_CONNECTION_TIMEOUT = 60
57 
58 # Maximum time to wait for a response message on the socket.
59 _SOCKET_READ_TIMEOUT = callback_handler.MAX_TIMEOUT
60 
61 
62 class Error(errors.DeviceError):
63     pass
64 
65 
66 class AppStartError(Error):
67     """Raised when the app is not able to be started."""
68 
69 
70 class AppRestoreConnectionError(Error):
71     """Raised when failed to restore app from disconnection."""
72 
73 
74 class ApiError(Error):
75     """Raised when remote API reports an error."""
76 
77 
78 class ProtocolError(Error):
79     """Raised when there is some error in exchanging data with server."""
80     NO_RESPONSE_FROM_HANDSHAKE = 'No response from handshake.'
81     NO_RESPONSE_FROM_SERVER = 'No response from server.'
82     MISMATCHED_API_ID = 'Mismatched API id.'
83 
84 
85 class JsonRpcCommand(object):
86     """Commands that can be invoked on all jsonrpc clients.
87 
88     INIT: Initializes a new session.
89     CONTINUE: Creates a connection.
90     """
91     INIT = 'initiate'
92     CONTINUE = 'continue'
93 
94 
95 class JsonRpcClientBase(object):
96     """Base class for jsonrpc clients that connect to remote servers.
97 
98     Connects to a remote device running a jsonrpc-compatible app. Before opening
99     a connection a port forward must be setup to go over usb. This be done using
100     adb.forward([local, remote]). Once the port has been forwarded it can be
101     used in this object as the port of communication.
102 
103     Attributes:
104         host_port: (int) The host port of this RPC client.
105         device_port: (int) The device port of this RPC client.
106         app_name: (str) The user-visible name of the app being communicated
107                   with.
108         uid: (int) The uid of this session.
109     """
110 
111     def __init__(self, app_name, ad):
112         """
113         Args:
114             app_name: (str) The user-visible name of the app being communicated
115                 with.
116             ad: (AndroidDevice) The device object associated with a client.
117         """
118         self.host_port = None
119         self.device_port = None
120         self.app_name = app_name
121         self._ad = ad
122         self.log = self._ad.log
123         self.uid = None
124         self._client = None  # prevent close errors on connect failure
125         self._conn = None
126         self._counter = None
127         self._lock = threading.Lock()
128         self._event_client = None
129 
130     def __del__(self):
131         self.disconnect()
132 
133     # Methods to be implemented by subclasses.
134 
135     def start_app_and_connect(self):
136         """Starts the server app on the android device and connects to it.
137 
138         After this, the self.host_port and self.device_port attributes must be
139         set.
140 
141         Must be implemented by subclasses.
142 
143         Raises:
144             AppStartError: When the app was not able to be started.
145         """
146         raise NotImplementedError()
147 
148     def stop_app(self):
149         """Kills any running instance of the app.
150 
151         Must be implemented by subclasses.
152         """
153         raise NotImplementedError()
154 
155     def restore_app_connection(self, port=None):
156         """Reconnects to the app after device USB was disconnected.
157 
158         Instead of creating new instance of the client:
159           - Uses the given port (or finds a new available host_port if none is
160             given).
161           - Tries to connect to remote server with selected port.
162 
163         Must be implemented by subclasses.
164 
165         Args:
166           port: If given, this is the host port from which to connect to remote
167               device port. If not provided, find a new available port as host
168               port.
169 
170         Raises:
171             AppRestoreConnectionError: When the app was not able to be
172             reconnected.
173         """
174         raise NotImplementedError()
175 
176     def _start_event_client(self):
177         """Starts a separate JsonRpc client to the same session for propagating
178         events.
179 
180         This is an optional function that should only implement if the client
181         utilizes the snippet event mechanism.
182 
183         Returns:
184             A JsonRpc Client object that connects to the same session as the
185             one on which this function is called.
186         """
187         raise NotImplementedError()
188 
189     # Rest of the client methods.
190 
191     def connect(self, uid=UNKNOWN_UID, cmd=JsonRpcCommand.INIT):
192         """Opens a connection to a JSON RPC server.
193 
194         Opens a connection to a remote client. The connection attempt will time
195         out if it takes longer than _SOCKET_CONNECTION_TIMEOUT seconds. Each
196         subsequent operation over this socket will time out after
197         _SOCKET_READ_TIMEOUT seconds as well.
198 
199         Args:
200             uid: int, The uid of the session to join, or UNKNOWN_UID to start a
201                 new session.
202             cmd: JsonRpcCommand, The command to use for creating the connection.
203 
204         Raises:
205             IOError: Raised when the socket times out from io error
206             socket.timeout: Raised when the socket waits to long for connection.
207             ProtocolError: Raised when there is an error in the protocol.
208         """
209         self._counter = self._id_counter()
210         self._conn = socket.create_connection(('localhost', self.host_port),
211                                               _SOCKET_CONNECTION_TIMEOUT)
212         self._conn.settimeout(_SOCKET_READ_TIMEOUT)
213         self._client = self._conn.makefile(mode='brw')
214 
215         resp = self._cmd(cmd, uid)
216         if not resp:
217             raise ProtocolError(
218                 self._ad, ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)
219         result = json.loads(str(resp, encoding='utf8'))
220         if result['status']:
221             self.uid = result['uid']
222         else:
223             self.uid = UNKNOWN_UID
224 
225     def disconnect(self):
226         """Close the connection to the remote client."""
227         if self._conn:
228             self._conn.close()
229             self._conn = None
230 
231     def _cmd(self, command, uid=None):
232         """Send a command to the server.
233 
234         Args:
235             command: str, The name of the command to execute.
236             uid: int, the uid of the session to send the command to.
237 
238         Returns:
239             The line that was written back.
240         """
241         if not uid:
242             uid = self.uid
243         self._client.write(
244             json.dumps({
245                 'cmd': command,
246                 'uid': uid
247             }).encode("utf8") + b'\n')
248         self._client.flush()
249         return self._client.readline()
250 
251     def _rpc(self, method, *args):
252         """Sends an rpc to the app.
253 
254         Args:
255             method: str, The name of the method to execute.
256             args: any, The args of the method.
257 
258         Returns:
259             The result of the rpc.
260 
261         Raises:
262             ProtocolError: Something went wrong with the protocol.
263             ApiError: The rpc went through, however executed with errors.
264         """
265         with self._lock:
266             apiid = next(self._counter)
267             data = {'id': apiid, 'method': method, 'params': args}
268             request = json.dumps(data)
269             self._client.write(request.encode("utf8") + b'\n')
270             self._client.flush()
271             response = self._client.readline()
272         if not response:
273             raise ProtocolError(self._ad,
274                                 ProtocolError.NO_RESPONSE_FROM_SERVER)
275         result = json.loads(str(response, encoding="utf8"))
276         if result['error']:
277             raise ApiError(self._ad, result['error'])
278         if result['id'] != apiid:
279             raise ProtocolError(self._ad, ProtocolError.MISMATCHED_API_ID)
280         if result.get('callback') is not None:
281             if self._event_client is None:
282                 self._event_client = self._start_event_client()
283             return callback_handler.CallbackHandler(
284                 callback_id=result['callback'],
285                 event_client=self._event_client,
286                 ret_value=result['result'],
287                 method_name=method,
288                 ad=self._ad)
289         return result['result']
290 
291     def __getattr__(self, name):
292         """Wrapper for python magic to turn method calls into RPC calls."""
293 
294         def rpc_call(*args):
295             return self._rpc(name, *args)
296 
297         return rpc_call
298 
299     def _id_counter(self):
300         i = 0
301         while True:
302             yield i
303             i += 1
304
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mobly/controllers/android_device_lib/jsonrpc_client_base.py b/mobly/controllers/android_device_lib/jsonrpc_client_base.py
--- a/mobly/controllers/android_device_lib/jsonrpc_client_base.py
+++ b/mobly/controllers/android_device_lib/jsonrpc_client_base.py
@@ -214,8 +214,8 @@
 
         resp = self._cmd(cmd, uid)
         if not resp:
-            raise ProtocolError(
-                self._ad, ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)
+            raise ProtocolError(self._ad,
+                                ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)
         result = json.loads(str(resp, encoding='utf8'))
         if result['status']:
             self.uid = result['uid']
@@ -228,6 +228,40 @@
             self._conn.close()
             self._conn = None
 
+    def _client_send(self, msg):
+        """Sends an Rpc message through the connection.
+
+        Args:
+            msg: string, the message to send.
+
+        Raises:
+            Error: a socket error occurred during the send.
+        """
+        try:
+            self._client.write(msg.encode("utf8") + b'\n')
+            self._client.flush()
+        except socket.error as e:
+            raise Error(
+                self._ad,
+                'Encountered socket error "%s" sending RPC message "%s"' %
+                (e, msg))
+
+    def _client_receive(self):
+        """Receives the server's response of an Rpc message.
+
+        Returns:
+            Raw byte string of the response.
+
+        Raises:
+            Error: a socket error occurred during the read.
+        """
+        try:
+            return self._client.readline()
+        except socket.error as e:
+            raise Error(
+                self._ad,
+                'Encountered socket error reading RPC response "%s"' % e)
+
     def _cmd(self, command, uid=None):
         """Send a command to the server.
 
@@ -240,13 +274,8 @@
         """
         if not uid:
             uid = self.uid
-        self._client.write(
-            json.dumps({
-                'cmd': command,
-                'uid': uid
-            }).encode("utf8") + b'\n')
-        self._client.flush()
-        return self._client.readline()
+        self._client_send(json.dumps({'cmd': command, 'uid': uid}))
+        return self._client_receive()
 
     def _rpc(self, method, *args):
         """Sends an rpc to the app.
@@ -266,9 +295,8 @@
             apiid = next(self._counter)
             data = {'id': apiid, 'method': method, 'params': args}
             request = json.dumps(data)
-            self._client.write(request.encode("utf8") + b'\n')
-            self._client.flush()
-            response = self._client.readline()
+            self._client_send(request)
+            response = self._client_receive()
         if not response:
             raise ProtocolError(self._ad,
                                 ProtocolError.NO_RESPONSE_FROM_SERVER)

{"golden_diff": "diff --git a/mobly/controllers/android_device_lib/jsonrpc_client_base.py b/mobly/controllers/android_device_lib/jsonrpc_client_base.py\n--- a/mobly/controllers/android_device_lib/jsonrpc_client_base.py\n+++ b/mobly/controllers/android_device_lib/jsonrpc_client_base.py\n@@ -214,8 +214,8 @@\n \n         resp = self._cmd(cmd, uid)\n         if not resp:\n-            raise ProtocolError(\n-                self._ad, ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)\n+            raise ProtocolError(self._ad,\n+                                ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)\n         result = json.loads(str(resp, encoding='utf8'))\n         if result['status']:\n             self.uid = result['uid']\n@@ -228,6 +228,40 @@\n             self._conn.close()\n             self._conn = None\n \n+    def _client_send(self, msg):\n+        \"\"\"Sends an Rpc message through the connection.\n+\n+        Args:\n+            msg: string, the message to send.\n+\n+        Raises:\n+            Error: a socket error occurred during the send.\n+        \"\"\"\n+        try:\n+            self._client.write(msg.encode(\"utf8\") + b'\\n')\n+            self._client.flush()\n+        except socket.error as e:\n+            raise Error(\n+                self._ad,\n+                'Encountered socket error \"%s\" sending RPC message \"%s\"' %\n+                (e, msg))\n+\n+    def _client_receive(self):\n+        \"\"\"Receives the server's response of an Rpc message.\n+\n+        Returns:\n+            Raw byte string of the response.\n+\n+        Raises:\n+            Error: a socket error occurred during the read.\n+        \"\"\"\n+        try:\n+            return self._client.readline()\n+        except socket.error as e:\n+            raise Error(\n+                self._ad,\n+                'Encountered socket error reading RPC response \"%s\"' % e)\n+\n     def _cmd(self, command, uid=None):\n         \"\"\"Send a command to the server.\n \n@@ -240,13 +274,8 @@\n         \"\"\"\n         if not uid:\n             uid = self.uid\n-        self._client.write(\n-            json.dumps({\n-                'cmd': command,\n-                'uid': uid\n-            }).encode(\"utf8\") + b'\\n')\n-        self._client.flush()\n-        return self._client.readline()\n+        self._client_send(json.dumps({'cmd': command, 'uid': uid}))\n+        return self._client_receive()\n \n     def _rpc(self, method, *args):\n         \"\"\"Sends an rpc to the app.\n@@ -266,9 +295,8 @@\n             apiid = next(self._counter)\n             data = {'id': apiid, 'method': method, 'params': args}\n             request = json.dumps(data)\n-            self._client.write(request.encode(\"utf8\") + b'\\n')\n-            self._client.flush()\n-            response = self._client.readline()\n+            self._client_send(request)\n+            response = self._client_receive()\n         if not response:\n             raise ProtocolError(self._ad,\n                                 ProtocolError.NO_RESPONSE_FROM_SERVER)\n", "issue": "Socket errors thrown from device obj are not marked with device tag\nWe should wrap `socket.error` with a child class of `DeviceError`.\n", "before_files": [{"content": "# Copyright 2016 Google Inc.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Base class for clients that communicate with apps over a JSON RPC interface.\n\nThe JSON protocol expected by this module is:\n\n.. code-block:: json\n\n    Request:\n    {\n        \"id\": <monotonically increasing integer containing the ID of \n               this request>\n        \"method\": <string containing the name of the method to execute>\n        \"params\": <JSON array containing the arguments to the method>\n    }\n\n    Response:\n    {\n        \"id\": <int id of request that this response maps to>,\n        \"result\": <Arbitrary JSON object containing the result of\n                   executing the method. If the method could not be\n                   executed or returned void, contains 'null'.>,\n        \"error\": <String containing the error thrown by executing the\n                  method. If no error occurred, contains 'null'.>\n        \"callback\": <String that represents a callback ID used to\n                     identify events associated with a particular\n                     CallbackHandler object.>\n    }\n\"\"\"\n\nfrom builtins import str\n\nimport json\nimport logging\nimport socket\nimport threading\n\nfrom mobly.controllers.android_device_lib import callback_handler\nfrom mobly.controllers.android_device_lib import errors\n\n# UID of the 'unknown' jsonrpc session. Will cause creation of a new session.\nUNKNOWN_UID = -1\n\n# Maximum time to wait for the socket to open on the device.\n_SOCKET_CONNECTION_TIMEOUT = 60\n\n# Maximum time to wait for a response message on the socket.\n_SOCKET_READ_TIMEOUT = callback_handler.MAX_TIMEOUT\n\n\nclass Error(errors.DeviceError):\n    pass\n\n\nclass AppStartError(Error):\n    \"\"\"Raised when the app is not able to be started.\"\"\"\n\n\nclass AppRestoreConnectionError(Error):\n    \"\"\"Raised when failed to restore app from disconnection.\"\"\"\n\n\nclass ApiError(Error):\n    \"\"\"Raised when remote API reports an error.\"\"\"\n\n\nclass ProtocolError(Error):\n    \"\"\"Raised when there is some error in exchanging data with server.\"\"\"\n    NO_RESPONSE_FROM_HANDSHAKE = 'No response from handshake.'\n    NO_RESPONSE_FROM_SERVER = 'No response from server.'\n    MISMATCHED_API_ID = 'Mismatched API id.'\n\n\nclass JsonRpcCommand(object):\n    \"\"\"Commands that can be invoked on all jsonrpc clients.\n\n    INIT: Initializes a new session.\n    CONTINUE: Creates a connection.\n    \"\"\"\n    INIT = 'initiate'\n    CONTINUE = 'continue'\n\n\nclass JsonRpcClientBase(object):\n    \"\"\"Base class for jsonrpc clients that connect to remote servers.\n\n    Connects to a remote device running a jsonrpc-compatible app. Before opening\n    a connection a port forward must be setup to go over usb. This be done using\n    adb.forward([local, remote]). Once the port has been forwarded it can be\n    used in this object as the port of communication.\n\n    Attributes:\n        host_port: (int) The host port of this RPC client.\n        device_port: (int) The device port of this RPC client.\n        app_name: (str) The user-visible name of the app being communicated\n                  with.\n        uid: (int) The uid of this session.\n    \"\"\"\n\n    def __init__(self, app_name, ad):\n        \"\"\"\n        Args:\n            app_name: (str) The user-visible name of the app being communicated\n                with.\n            ad: (AndroidDevice) The device object associated with a client.\n        \"\"\"\n        self.host_port = None\n        self.device_port = None\n        self.app_name = app_name\n        self._ad = ad\n        self.log = self._ad.log\n        self.uid = None\n        self._client = None  # prevent close errors on connect failure\n        self._conn = None\n        self._counter = None\n        self._lock = threading.Lock()\n        self._event_client = None\n\n    def __del__(self):\n        self.disconnect()\n\n    # Methods to be implemented by subclasses.\n\n    def start_app_and_connect(self):\n        \"\"\"Starts the server app on the android device and connects to it.\n\n        After this, the self.host_port and self.device_port attributes must be\n        set.\n\n        Must be implemented by subclasses.\n\n        Raises:\n            AppStartError: When the app was not able to be started.\n        \"\"\"\n        raise NotImplementedError()\n\n    def stop_app(self):\n        \"\"\"Kills any running instance of the app.\n\n        Must be implemented by subclasses.\n        \"\"\"\n        raise NotImplementedError()\n\n    def restore_app_connection(self, port=None):\n        \"\"\"Reconnects to the app after device USB was disconnected.\n\n        Instead of creating new instance of the client:\n          - Uses the given port (or finds a new available host_port if none is\n            given).\n          - Tries to connect to remote server with selected port.\n\n        Must be implemented by subclasses.\n\n        Args:\n          port: If given, this is the host port from which to connect to remote\n              device port. If not provided, find a new available port as host\n              port.\n\n        Raises:\n            AppRestoreConnectionError: When the app was not able to be\n            reconnected.\n        \"\"\"\n        raise NotImplementedError()\n\n    def _start_event_client(self):\n        \"\"\"Starts a separate JsonRpc client to the same session for propagating\n        events.\n\n        This is an optional function that should only implement if the client\n        utilizes the snippet event mechanism.\n\n        Returns:\n            A JsonRpc Client object that connects to the same session as the\n            one on which this function is called.\n        \"\"\"\n        raise NotImplementedError()\n\n    # Rest of the client methods.\n\n    def connect(self, uid=UNKNOWN_UID, cmd=JsonRpcCommand.INIT):\n        \"\"\"Opens a connection to a JSON RPC server.\n\n        Opens a connection to a remote client. The connection attempt will time\n        out if it takes longer than _SOCKET_CONNECTION_TIMEOUT seconds. Each\n        subsequent operation over this socket will time out after\n        _SOCKET_READ_TIMEOUT seconds as well.\n\n        Args:\n            uid: int, The uid of the session to join, or UNKNOWN_UID to start a\n                new session.\n            cmd: JsonRpcCommand, The command to use for creating the connection.\n\n        Raises:\n            IOError: Raised when the socket times out from io error\n            socket.timeout: Raised when the socket waits to long for connection.\n            ProtocolError: Raised when there is an error in the protocol.\n        \"\"\"\n        self._counter = self._id_counter()\n        self._conn = socket.create_connection(('localhost', self.host_port),\n                                              _SOCKET_CONNECTION_TIMEOUT)\n        self._conn.settimeout(_SOCKET_READ_TIMEOUT)\n        self._client = self._conn.makefile(mode='brw')\n\n        resp = self._cmd(cmd, uid)\n        if not resp:\n            raise ProtocolError(\n                self._ad, ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)\n        result = json.loads(str(resp, encoding='utf8'))\n        if result['status']:\n            self.uid = result['uid']\n        else:\n            self.uid = UNKNOWN_UID\n\n    def disconnect(self):\n        \"\"\"Close the connection to the remote client.\"\"\"\n        if self._conn:\n            self._conn.close()\n            self._conn = None\n\n    def _cmd(self, command, uid=None):\n        \"\"\"Send a command to the server.\n\n        Args:\n            command: str, The name of the command to execute.\n            uid: int, the uid of the session to send the command to.\n\n        Returns:\n            The line that was written back.\n        \"\"\"\n        if not uid:\n            uid = self.uid\n        self._client.write(\n            json.dumps({\n                'cmd': command,\n                'uid': uid\n            }).encode(\"utf8\") + b'\\n')\n        self._client.flush()\n        return self._client.readline()\n\n    def _rpc(self, method, *args):\n        \"\"\"Sends an rpc to the app.\n\n        Args:\n            method: str, The name of the method to execute.\n            args: any, The args of the method.\n\n        Returns:\n            The result of the rpc.\n\n        Raises:\n            ProtocolError: Something went wrong with the protocol.\n            ApiError: The rpc went through, however executed with errors.\n        \"\"\"\n        with self._lock:\n            apiid = next(self._counter)\n            data = {'id': apiid, 'method': method, 'params': args}\n            request = json.dumps(data)\n            self._client.write(request.encode(\"utf8\") + b'\\n')\n            self._client.flush()\n            response = self._client.readline()\n        if not response:\n            raise ProtocolError(self._ad,\n                                ProtocolError.NO_RESPONSE_FROM_SERVER)\n        result = json.loads(str(response, encoding=\"utf8\"))\n        if result['error']:\n            raise ApiError(self._ad, result['error'])\n        if result['id'] != apiid:\n            raise ProtocolError(self._ad, ProtocolError.MISMATCHED_API_ID)\n        if result.get('callback') is not None:\n            if self._event_client is None:\n                self._event_client = self._start_event_client()\n            return callback_handler.CallbackHandler(\n                callback_id=result['callback'],\n                event_client=self._event_client,\n                ret_value=result['result'],\n                method_name=method,\n                ad=self._ad)\n        return result['result']\n\n    def __getattr__(self, name):\n        \"\"\"Wrapper for python magic to turn method calls into RPC calls.\"\"\"\n\n        def rpc_call(*args):\n            return self._rpc(name, *args)\n\n        return rpc_call\n\n    def _id_counter(self):\n        i = 0\n        while True:\n            yield i\n            i += 1\n", "path": "mobly/controllers/android_device_lib/jsonrpc_client_base.py"}], "after_files": [{"content": "# Copyright 2016 Google Inc.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Base class for clients that communicate with apps over a JSON RPC interface.\n\nThe JSON protocol expected by this module is:\n\n.. code-block:: json\n\n    Request:\n    {\n        \"id\": <monotonically increasing integer containing the ID of \n               this request>\n        \"method\": <string containing the name of the method to execute>\n        \"params\": <JSON array containing the arguments to the method>\n    }\n\n    Response:\n    {\n        \"id\": <int id of request that this response maps to>,\n        \"result\": <Arbitrary JSON object containing the result of\n                   executing the method. If the method could not be\n                   executed or returned void, contains 'null'.>,\n        \"error\": <String containing the error thrown by executing the\n                  method. If no error occurred, contains 'null'.>\n        \"callback\": <String that represents a callback ID used to\n                     identify events associated with a particular\n                     CallbackHandler object.>\n    }\n\"\"\"\n\nfrom builtins import str\n\nimport json\nimport logging\nimport socket\nimport threading\n\nfrom mobly.controllers.android_device_lib import callback_handler\nfrom mobly.controllers.android_device_lib import errors\n\n# UID of the 'unknown' jsonrpc session. Will cause creation of a new session.\nUNKNOWN_UID = -1\n\n# Maximum time to wait for the socket to open on the device.\n_SOCKET_CONNECTION_TIMEOUT = 60\n\n# Maximum time to wait for a response message on the socket.\n_SOCKET_READ_TIMEOUT = callback_handler.MAX_TIMEOUT\n\n\nclass Error(errors.DeviceError):\n    pass\n\n\nclass AppStartError(Error):\n    \"\"\"Raised when the app is not able to be started.\"\"\"\n\n\nclass AppRestoreConnectionError(Error):\n    \"\"\"Raised when failed to restore app from disconnection.\"\"\"\n\n\nclass ApiError(Error):\n    \"\"\"Raised when remote API reports an error.\"\"\"\n\n\nclass ProtocolError(Error):\n    \"\"\"Raised when there is some error in exchanging data with server.\"\"\"\n    NO_RESPONSE_FROM_HANDSHAKE = 'No response from handshake.'\n    NO_RESPONSE_FROM_SERVER = 'No response from server.'\n    MISMATCHED_API_ID = 'Mismatched API id.'\n\n\nclass JsonRpcCommand(object):\n    \"\"\"Commands that can be invoked on all jsonrpc clients.\n\n    INIT: Initializes a new session.\n    CONTINUE: Creates a connection.\n    \"\"\"\n    INIT = 'initiate'\n    CONTINUE = 'continue'\n\n\nclass JsonRpcClientBase(object):\n    \"\"\"Base class for jsonrpc clients that connect to remote servers.\n\n    Connects to a remote device running a jsonrpc-compatible app. Before opening\n    a connection a port forward must be setup to go over usb. This be done using\n    adb.forward([local, remote]). Once the port has been forwarded it can be\n    used in this object as the port of communication.\n\n    Attributes:\n        host_port: (int) The host port of this RPC client.\n        device_port: (int) The device port of this RPC client.\n        app_name: (str) The user-visible name of the app being communicated\n                  with.\n        uid: (int) The uid of this session.\n    \"\"\"\n\n    def __init__(self, app_name, ad):\n        \"\"\"\n        Args:\n            app_name: (str) The user-visible name of the app being communicated\n                with.\n            ad: (AndroidDevice) The device object associated with a client.\n        \"\"\"\n        self.host_port = None\n        self.device_port = None\n        self.app_name = app_name\n        self._ad = ad\n        self.log = self._ad.log\n        self.uid = None\n        self._client = None  # prevent close errors on connect failure\n        self._conn = None\n        self._counter = None\n        self._lock = threading.Lock()\n        self._event_client = None\n\n    def __del__(self):\n        self.disconnect()\n\n    # Methods to be implemented by subclasses.\n\n    def start_app_and_connect(self):\n        \"\"\"Starts the server app on the android device and connects to it.\n\n        After this, the self.host_port and self.device_port attributes must be\n        set.\n\n        Must be implemented by subclasses.\n\n        Raises:\n            AppStartError: When the app was not able to be started.\n        \"\"\"\n        raise NotImplementedError()\n\n    def stop_app(self):\n        \"\"\"Kills any running instance of the app.\n\n        Must be implemented by subclasses.\n        \"\"\"\n        raise NotImplementedError()\n\n    def restore_app_connection(self, port=None):\n        \"\"\"Reconnects to the app after device USB was disconnected.\n\n        Instead of creating new instance of the client:\n          - Uses the given port (or finds a new available host_port if none is\n            given).\n          - Tries to connect to remote server with selected port.\n\n        Must be implemented by subclasses.\n\n        Args:\n          port: If given, this is the host port from which to connect to remote\n              device port. If not provided, find a new available port as host\n              port.\n\n        Raises:\n            AppRestoreConnectionError: When the app was not able to be\n            reconnected.\n        \"\"\"\n        raise NotImplementedError()\n\n    def _start_event_client(self):\n        \"\"\"Starts a separate JsonRpc client to the same session for propagating\n        events.\n\n        This is an optional function that should only implement if the client\n        utilizes the snippet event mechanism.\n\n        Returns:\n            A JsonRpc Client object that connects to the same session as the\n            one on which this function is called.\n        \"\"\"\n        raise NotImplementedError()\n\n    # Rest of the client methods.\n\n    def connect(self, uid=UNKNOWN_UID, cmd=JsonRpcCommand.INIT):\n        \"\"\"Opens a connection to a JSON RPC server.\n\n        Opens a connection to a remote client. The connection attempt will time\n        out if it takes longer than _SOCKET_CONNECTION_TIMEOUT seconds. Each\n        subsequent operation over this socket will time out after\n        _SOCKET_READ_TIMEOUT seconds as well.\n\n        Args:\n            uid: int, The uid of the session to join, or UNKNOWN_UID to start a\n                new session.\n            cmd: JsonRpcCommand, The command to use for creating the connection.\n\n        Raises:\n            IOError: Raised when the socket times out from io error\n            socket.timeout: Raised when the socket waits to long for connection.\n            ProtocolError: Raised when there is an error in the protocol.\n        \"\"\"\n        self._counter = self._id_counter()\n        self._conn = socket.create_connection(('localhost', self.host_port),\n                                              _SOCKET_CONNECTION_TIMEOUT)\n        self._conn.settimeout(_SOCKET_READ_TIMEOUT)\n        self._client = self._conn.makefile(mode='brw')\n\n        resp = self._cmd(cmd, uid)\n        if not resp:\n            raise ProtocolError(self._ad,\n                                ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)\n        result = json.loads(str(resp, encoding='utf8'))\n        if result['status']:\n            self.uid = result['uid']\n        else:\n            self.uid = UNKNOWN_UID\n\n    def disconnect(self):\n        \"\"\"Close the connection to the remote client.\"\"\"\n        if self._conn:\n            self._conn.close()\n            self._conn = None\n\n    def _client_send(self, msg):\n        \"\"\"Sends an Rpc message through the connection.\n\n        Args:\n            msg: string, the message to send.\n\n        Raises:\n            Error: a socket error occurred during the send.\n        \"\"\"\n        try:\n            self._client.write(msg.encode(\"utf8\") + b'\\n')\n            self._client.flush()\n        except socket.error as e:\n            raise Error(\n                self._ad,\n                'Encountered socket error \"%s\" sending RPC message \"%s\"' %\n                (e, msg))\n\n    def _client_receive(self):\n        \"\"\"Receives the server's response of an Rpc message.\n\n        Returns:\n            Raw byte string of the response.\n\n        Raises:\n            Error: a socket error occurred during the read.\n        \"\"\"\n        try:\n            return self._client.readline()\n        except socket.error as e:\n            raise Error(\n                self._ad,\n                'Encountered socket error reading RPC response \"%s\"' % e)\n\n    def _cmd(self, command, uid=None):\n        \"\"\"Send a command to the server.\n\n        Args:\n            command: str, The name of the command to execute.\n            uid: int, the uid of the session to send the command to.\n\n        Returns:\n            The line that was written back.\n        \"\"\"\n        if not uid:\n            uid = self.uid\n        self._client_send(json.dumps({'cmd': command, 'uid': uid}))\n        return self._client_receive()\n\n    def _rpc(self, method, *args):\n        \"\"\"Sends an rpc to the app.\n\n        Args:\n            method: str, The name of the method to execute.\n            args: any, The args of the method.\n\n        Returns:\n            The result of the rpc.\n\n        Raises:\n            ProtocolError: Something went wrong with the protocol.\n            ApiError: The rpc went through, however executed with errors.\n        \"\"\"\n        with self._lock:\n            apiid = next(self._counter)\n            data = {'id': apiid, 'method': method, 'params': args}\n            request = json.dumps(data)\n            self._client_send(request)\n            response = self._client_receive()\n        if not response:\n            raise ProtocolError(self._ad,\n                                ProtocolError.NO_RESPONSE_FROM_SERVER)\n        result = json.loads(str(response, encoding=\"utf8\"))\n        if result['error']:\n            raise ApiError(self._ad, result['error'])\n        if result['id'] != apiid:\n            raise ProtocolError(self._ad, ProtocolError.MISMATCHED_API_ID)\n        if result.get('callback') is not None:\n            if self._event_client is None:\n                self._event_client = self._start_event_client()\n            return callback_handler.CallbackHandler(\n                callback_id=result['callback'],\n                event_client=self._event_client,\n                ret_value=result['result'],\n                method_name=method,\n                ad=self._ad)\n        return result['result']\n\n    def __getattr__(self, name):\n        \"\"\"Wrapper for python magic to turn method calls into RPC calls.\"\"\"\n\n        def rpc_call(*args):\n            return self._rpc(name, *args)\n\n        return rpc_call\n\n    def _id_counter(self):\n        i = 0\n        while True:\n            yield i\n            i += 1\n", "path": "mobly/controllers/android_device_lib/jsonrpc_client_base.py"}]}