xref: /aosp_15_r20/tools/acloud/internal/lib/emulator_console_test.py (revision 800a58d989c669b8eb8a71d8df53b1ba3d411444) 
1#!/usr/bin/env python3
2#
3# Copyright 2021 - 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"""Unit tests for RemoteEmulatorConsole."""
18
19import unittest
20import subprocess
21
22from unittest import mock
23
24from acloud import errors
25from acloud.internal.lib import driver_test_lib
26from acloud.internal.lib import emulator_console
27
28
29class EmulatorConsoleTest(driver_test_lib.BaseDriverTest):
30    """Unit tests for RemoteEmulatorConsole."""
31
32    _LOCAL_PORT = 56789
33    _TIMEOUT_SECS = 100
34
35    def setUp(self):
36        """Create mock objects."""
37        super().setUp()
38        self._mock_pick_free_port = self.Patch(
39            emulator_console.utils, "PickFreePort",
40            return_value=self._LOCAL_PORT)
41        self._mock_establish_ssh_tunnel = self.Patch(
42            emulator_console.utils,
43            "EstablishSshTunnel")
44        self._mock_release_port = self.Patch(
45            emulator_console.utils,
46            "ReleasePort")
47        self._mock_connection = mock.Mock()
48        self._mock_create_connection = self.Patch(
49            emulator_console.socket,
50            "create_connection",
51            return_value=self._mock_connection)
52
53    def _CreateRemoteEmulatorConsole(self):
54        """_Create a RemoteEmulatorConsole."""
55        console = emulator_console.RemoteEmulatorConsole(
56            "192.0.2.1", 5444, "user", "key_path", "extra args",
57            self._TIMEOUT_SECS)
58        self._mock_pick_free_port.assert_called_once()
59        self._mock_establish_ssh_tunnel.assert_called_once_with(
60            "192.0.2.1", "key_path", "user",
61            [(self._LOCAL_PORT, 5444)], "extra args")
62        self._mock_create_connection.assert_called_once_with(
63            ("127.0.0.1", self._LOCAL_PORT), self._TIMEOUT_SECS)
64        self._mock_connection.settimeout.assert_called_once_with(
65            self._TIMEOUT_SECS)
66        self._mock_release_port.assert_not_called()
67        self._mock_connection.close.assert_not_called()
68
69        self._mock_pick_free_port.reset_mock()
70        self._mock_establish_ssh_tunnel.reset_mock()
71        self._mock_create_connection.reset_mock()
72        self._mock_connection.settimeout.reset_mock()
73        return console
74
75    def testInitSshTunnelError(self):
76        """Test not releasing port if SSH tunnel fails."""
77        self._mock_establish_ssh_tunnel.side_effect = (
78            subprocess.CalledProcessError(returncode=1, cmd="ssh"))
79        with self.assertRaises(errors.DeviceConnectionError):
80            emulator_console.RemoteEmulatorConsole(
81                "192.0.2.1", 5444, "user", "key_path", "extra args",
82                self._TIMEOUT_SECS)
83        self._mock_connection.settimeout.assert_not_called()
84        self._mock_connection.close.assert_not_called()
85        self._mock_release_port.assert_not_called()
86
87    def testInitConnectionError(self):
88        """Test releasing port when create_connection fails."""
89        self._mock_create_connection.side_effect = OSError()
90        with self.assertRaises(errors.DeviceConnectionError):
91            emulator_console.RemoteEmulatorConsole(
92                "192.0.2.1", 5444, "user", "key_path", "extra args",
93                self._TIMEOUT_SECS)
94        self._mock_connection.settimeout.assert_not_called()
95        self._mock_connection.close.assert_not_called()
96        self._mock_release_port.assert_called_once()
97
98    def testInitSocketError(self):
99        """Test closing socket when settimeout fails."""
100        self._mock_connection.settimeout.side_effect = OSError()
101        with self.assertRaises(errors.DeviceConnectionError):
102            emulator_console.RemoteEmulatorConsole(
103                "192.0.2.1", 5444, "user", "key_path", "extra args",
104                self._TIMEOUT_SECS)
105        self._mock_connection.settimeout.assert_called_once()
106        self._mock_connection.close.assert_called_once()
107        self._mock_release_port.assert_called_once()
108
109    def testContext(self):
110        """Test RemoteEmulatorConsole as a context manager."""
111        with self._CreateRemoteEmulatorConsole():
112            pass
113        self._mock_connection.close.assert_called_once()
114        self._mock_release_port.assert_called_once()
115
116    def testReconnect(self):
117        """Test RemoteEmulatorConsole.Reconnect."""
118        console = self._CreateRemoteEmulatorConsole()
119        console.Reconnect()
120        self._mock_pick_free_port.assert_not_called()
121        self._mock_establish_ssh_tunnel.assert_not_called()
122        self._mock_release_port.assert_not_called()
123        self._mock_connection.close.assert_called_once()
124        self._mock_create_connection.assert_called_once_with(
125            ("127.0.0.1", self._LOCAL_PORT), self._TIMEOUT_SECS)
126        self._mock_connection.settimeout.assert_called_once_with(
127            self._TIMEOUT_SECS)
128
129    def testSend(self):
130        """Test RemoteEmulatorConsole.Send."""
131        console = self._CreateRemoteEmulatorConsole()
132        console.Send("ping")
133        self._mock_connection.send.assert_called_with(b"ping\n")
134
135        self._mock_connection.send.side_effect = OSError()
136        with self.assertRaises(errors.DeviceConnectionError):
137            console.Send("ping")
138
139    def testRecv(self):
140        """Test RemoteEmulatorConsole.Recv."""
141        console = self._CreateRemoteEmulatorConsole()
142
143        self._mock_connection.recv.side_effect = [b"1", b"1"]
144        self.assertEqual("11", console.Recv("11", buffer_size=1))
145        self._mock_connection.recv.side_effect = [b"12"]
146        self.assertEqual("12", console.Recv("2"))
147
148        self._mock_connection.recv.side_effect = [b"1", b""]
149        with self.assertRaises(errors.DeviceConnectionError):
150            console.Recv("2")
151        self._mock_connection.recv.side_effect = OSError()
152        with self.assertRaises(errors.DeviceConnectionError):
153            console.Recv("1")
154
155    def testPing(self):
156        """Test RemoteEmulatorConsole.Ping."""
157        console = self._CreateRemoteEmulatorConsole()
158
159        self._mock_connection.recv.side_effect = [b"I am alive!\r\nOK\r\n"]
160        self.assertTrue(console.Ping())
161
162        self._mock_connection.recv.side_effect = [b""]
163        self.assertFalse(console.Ping())
164
165        self._mock_connection.recv.side_effect = OSError()
166        self.assertFalse(console.Ping())
167
168    def testKill(self):
169        """Test RemoteEmulatorConsole.Kill."""
170        console = self._CreateRemoteEmulatorConsole()
171        self._mock_connection.recv.side_effect = [b"bye bye\r\n"]
172        console.Kill()
173
174
175if __name__ == "__main__":
176    unittest.main()
177