1#!/usr/bin/env python
2#
3# Copyright (C) 2017 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import queue
19import socket
20import threading
21import unittest
22
23from host_controller.tradefed import remote_client
24from host_controller.tradefed import remote_operation
25
26
27class MockRemoteManagerThread(threading.Thread):
28    """A thread which mocks remote manager.
29
30    Attributes:
31        HOST: Local host name.
32        PORT: The port that the remote manager listens to.
33        _remote_mgr_socket: The remote manager socket.
34        _response_queue: A queue.Queue object containing the response strings.
35        _timeout: Socket timeout in seconds.
36        last_error: The exception which caused this thread to terminate.
37    """
38    HOST = remote_client.LOCALHOST
39    PORT = 32123
40
41    def __init__(self, timeout):
42        """Creates and listens to remote manager socket."""
43        super(MockRemoteManagerThread, self).__init__()
44        self._response_queue = queue.Queue()
45        self._timeout = timeout
46        self.last_error = None
47        self._remote_mgr_socket = socket.socket()
48        try:
49            self._remote_mgr_socket.settimeout(self._timeout)
50            self._remote_mgr_socket.bind((self.HOST, self.PORT))
51            self._remote_mgr_socket.listen(1)
52        except socket.error:
53            self._remote_mgr_socket.close()
54
55    def _Respond(self, response_str):
56        """Accepts a client connection and responds.
57
58        Args:
59            response_str: The response string.
60        """
61        (server_socket, client_address) = self._remote_mgr_socket.accept()
62        try:
63            server_socket.settimeout(self._timeout)
64            # Receive until connection is closed
65            while not server_socket.recv(4096):
66                pass
67            server_socket.send(response_str)
68        finally:
69            server_socket.close()
70
71    def AddResponse(self, response_str):
72        """Add a response string to the queue.
73
74        Args:
75            response_str: The response string.
76        """
77        self._response_queue.put_nowait(response_str)
78
79    def CloseSocket(self):
80        """Closes the remote manager socket."""
81        if self._remote_mgr_socket:
82            self._remote_mgr_socket.close()
83            self._remote_mgr_socket = None
84
85    # @Override
86    def run(self):
87        """Sends the queued responses to the clients."""
88        try:
89            while True:
90                response_str = self._response_queue.get()
91                self._response_queue.task_done()
92                if response_str is None:
93                    break
94                self._Respond(response_str)
95        except socket.error as e:
96            self.last_error = e
97        finally:
98            self.CloseSocket()
99
100
101class RemoteClientTest(unittest.TestCase):
102    """A test for remote_client.RemoteClient.
103
104    Attributes:
105        _remote_mgr_thread: An instance of MockRemoteManagerThread.
106        _client: The remote_client.RemoteClient being tested.
107    """
108
109    def setUp(self):
110        """Creates remote manager thread."""
111        self._remote_mgr_thread = MockRemoteManagerThread(5)
112        self._remote_mgr_thread.daemon = True
113        self._remote_mgr_thread.start()
114        self._client = remote_client.RemoteClient(self._remote_mgr_thread.HOST,
115                                                  self._remote_mgr_thread.PORT,
116                                                  5)
117
118    def tearDown(self):
119        """Terminates remote manager thread."""
120        self._remote_mgr_thread.AddResponse(None)
121        self._remote_mgr_thread.join(15)
122        self._remote_mgr_thread.CloseSocket()
123        self.assertFalse(self._remote_mgr_thread.is_alive(),
124                         "Cannot stop remote manager thread.")
125        if self._remote_mgr_thread.last_error:
126            raise self._remote_mgr_thread.last_error
127
128    def testListDevice(self):
129        """Tests ListDevices operation."""
130        self._remote_mgr_thread.AddResponse('{"serials": []}')
131        self._client.ListDevices()
132
133    def testAddCommand(self):
134        """Tests AddCommand operation."""
135        self._remote_mgr_thread.AddResponse('{}')
136        self._client.SendOperation(remote_operation.AddCommand(0, "COMMAND"))
137
138    def testMultipleOperations(self):
139        """Tests sending multiple operations via one connection."""
140        self._remote_mgr_thread.AddResponse('{}\n{}')
141        self._client.SendOperations(remote_operation.ListDevices(),
142                                    remote_operation.ListDevices())
143
144    def testExecuteCommand(self):
145        """Tests executing a command and waiting for result."""
146        self._remote_mgr_thread.AddResponse('{}')
147        self._client.SendOperation(remote_operation.AllocateDevice("serial123"))
148        self._remote_mgr_thread.AddResponse('{}')
149        self._client.SendOperation(remote_operation.ExecuteCommand(
150                "serial123", "vts", "-m", "SampleShellTest"))
151
152        self._remote_mgr_thread.AddResponse('{"status": "EXECUTING"}')
153        result = self._client.WaitForCommandResult("serial123",
154                                                   timeout=0.5, poll_interval=1)
155        self.assertIsNone(result, "Client returns result before command finishes.")
156
157        self._remote_mgr_thread.AddResponse('{"status": "EXECUTING"}')
158        self._remote_mgr_thread.AddResponse('{"status": "INVOCATION_SUCCESS"}')
159        result = self._client.WaitForCommandResult("serial123",
160                                                   timeout=5, poll_interval=1)
161        self._remote_mgr_thread.AddResponse('{}')
162        self._client.SendOperation(remote_operation.FreeDevice("serial123"))
163        self.assertIsNotNone(result, "Client doesn't return command result.")
164
165    def testSocketError(self):
166        """Tests raising exception when socket error occurs."""
167        self.assertRaises(socket.timeout, self._client.ListDevices)
168        self._remote_mgr_thread.AddResponse(None)
169        self.assertRaises(socket.error, self._client.ListDevices)
170
171    def testRemoteOperationException(self):
172        """Tests raising exception when response is an error."""
173        self._remote_mgr_thread.AddResponse('{"error": "unit test"}')
174        self.assertRaises(remote_operation.RemoteOperationException,
175                          self._client.ListDevices)
176
177
178if __name__ == "__main__":
179    unittest.main()
180