1# Copyright 2018 - The Android Open Source Project
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"""Tests for reconnect."""
15
16import collections
17import unittest
18import subprocess
19
20import mock
21
22from acloud import errors
23from acloud.internal import constants
24from acloud.internal.lib import driver_test_lib
25from acloud.internal.lib import utils
26from acloud.internal.lib.adb_tools import AdbTools
27from acloud.reconnect import reconnect
28
29
30ForwardedPorts = collections.namedtuple("ForwardedPorts",
31                                        [constants.VNC_PORT, constants.ADB_PORT])
32
33
34class ReconnectTest(driver_test_lib.BaseDriverTest):
35    """Test reconnect functions."""
36
37    # pylint: disable=no-member, too-many-statements
38    def testReconnectInstance(self):
39        """Test Reconnect Instances."""
40        ssh_private_key_path = "/fake/acloud_rsa"
41        fake_report = mock.MagicMock()
42        instance_object = mock.MagicMock()
43        instance_object.name = "fake_name"
44        instance_object.ip = "1.1.1.1"
45        instance_object.islocal = False
46        instance_object.adb_port = "8686"
47        instance_object.avd_type = "cuttlefish"
48        self.Patch(subprocess, "check_call", return_value=True)
49        self.Patch(utils, "LaunchVncClient")
50        self.Patch(utils, "AutoConnect")
51        self.Patch(AdbTools, "IsAdbConnected", return_value=False)
52        self.Patch(AdbTools, "IsAdbConnectionAlive", return_value=False)
53        self.Patch(utils, "IsCommandRunning", return_value=False)
54        self.Patch(reconnect, "_IsWebrtcEnable", return_value=False)
55        fake_device_dict = {
56            constants.IP: "1.1.1.1",
57            constants.INSTANCE_NAME: "fake_name",
58            constants.VNC_PORT: 6666,
59            constants.ADB_PORT: "8686"
60        }
61
62        # test ssh tunnel not connected, remote instance.
63        instance_object.vnc_port = 6666
64        instance_object.display = ""
65        utils.AutoConnect.call_count = 0
66        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report)
67        utils.AutoConnect.assert_not_called()
68        utils.LaunchVncClient.assert_called_with(6666)
69        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)
70
71        instance_object.display = "888x777 (99)"
72        utils.AutoConnect.call_count = 0
73        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report)
74        utils.AutoConnect.assert_not_called()
75        utils.LaunchVncClient.assert_called_with(6666, "888", "777")
76        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)
77
78        # test ssh tunnel connected , remote instance.
79        instance_object.ssh_tunnel_is_connected = False
80        instance_object.display = ""
81        utils.AutoConnect.call_count = 0
82        instance_object.vnc_port = 5555
83        extra_args_ssh_tunnel = None
84        self.Patch(utils, "AutoConnect",
85                   return_value=ForwardedPorts(vnc_port=11111, adb_port=22222))
86        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report)
87        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
88                                             rsa_key_file=ssh_private_key_path,
89                                             target_vnc_port=constants.CF_VNC_PORT,
90                                             target_adb_port=constants.CF_ADB_PORT,
91                                             ssh_user=constants.GCE_USER,
92                                             extra_args_ssh_tunnel=extra_args_ssh_tunnel)
93        utils.LaunchVncClient.assert_called_with(11111)
94        fake_device_dict = {
95            constants.IP: "1.1.1.1",
96            constants.INSTANCE_NAME: "fake_name",
97            constants.VNC_PORT: 11111,
98            constants.ADB_PORT: 22222
99        }
100        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)
101
102        instance_object.display = "999x777 (99)"
103        extra_args_ssh_tunnel = "fake_extra_args_ssh_tunnel"
104        utils.AutoConnect.call_count = 0
105        reconnect.ReconnectInstance(ssh_private_key_path,
106                                    instance_object,
107                                    fake_report,
108                                    extra_args_ssh_tunnel)
109        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
110                                             rsa_key_file=ssh_private_key_path,
111                                             target_vnc_port=constants.CF_VNC_PORT,
112                                             target_adb_port=constants.CF_ADB_PORT,
113                                             ssh_user=constants.GCE_USER,
114                                             extra_args_ssh_tunnel=extra_args_ssh_tunnel)
115        utils.LaunchVncClient.assert_called_with(11111, "999", "777")
116        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)
117
118        # test fail reconnect report.
119        self.Patch(utils, "AutoConnect",
120                   return_value=ForwardedPorts(vnc_port=None, adb_port=None))
121        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report)
122        fake_device_dict = {
123            constants.IP: "1.1.1.1",
124            constants.INSTANCE_NAME: "fake_name",
125            constants.VNC_PORT: None,
126            constants.ADB_PORT: None
127        }
128        fake_report.AddData.assert_called_with(key="device_failing_reconnect",
129                                               value=fake_device_dict)
130
131        # test reconnect local instance.
132        instance_object.islocal = True
133        instance_object.display = ""
134        instance_object.vnc_port = 5555
135        instance_object.ssh_tunnel_is_connected = False
136        utils.AutoConnect.call_count = 0
137        reconnect.ReconnectInstance(ssh_private_key_path,
138                                    instance_object,
139                                    fake_report)
140        utils.AutoConnect.assert_not_called()
141        utils.LaunchVncClient.assert_called_with(5555)
142        fake_device_dict = {
143            constants.IP: "1.1.1.1",
144            constants.INSTANCE_NAME: "fake_name",
145            constants.VNC_PORT: 5555,
146            constants.ADB_PORT: "8686"
147        }
148        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)
149
150    # pylint: disable=no-member
151    def testReconnectInstanceWithWebRTC(self):
152        """Test reconnect instances with WebRTC."""
153        ssh_private_key_path = "/fake/acloud_rsa"
154        fake_report = mock.MagicMock()
155        instance_object = mock.MagicMock()
156        instance_object.ip = "1.1.1.1"
157        instance_object.islocal = False
158        instance_object.adb_port = "8686"
159        instance_object.avd_type = "cuttlefish"
160        self.Patch(subprocess, "check_call", return_value=True)
161        self.Patch(utils, "LaunchVncClient")
162        self.Patch(utils, "AutoConnect")
163        self.Patch(utils, "LaunchBrowser")
164        self.Patch(utils, "EstablishWebRTCSshTunnel")
165        self.Patch(AdbTools, "IsAdbConnected", return_value=False)
166        self.Patch(AdbTools, "IsAdbConnectionAlive", return_value=False)
167        self.Patch(utils, "IsCommandRunning", return_value=False)
168        self.Patch(reconnect, "_IsWebrtcEnable", return_value=True)
169
170        # test ssh tunnel not reconnect to the remote instance.
171        instance_object.vnc_port = 6666
172        instance_object.display = ""
173        utils.AutoConnect.call_count = 0
174        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report)
175        utils.AutoConnect.assert_not_called()
176        utils.LaunchVncClient.assert_not_called()
177        utils.EstablishWebRTCSshTunnel.assert_called_with(extra_args_ssh_tunnel=None,
178                                                          ip_addr='1.1.1.1',
179                                                          rsa_key_file='/fake/acloud_rsa',
180                                                          ssh_user='vsoc-01')
181        utils.LaunchBrowser.assert_called_with('localhost', 8443)
182
183    def testReconnectInstanceAvdtype(self):
184        """Test Reconnect Instances of avd_type."""
185        ssh_private_key_path = "/fake/acloud_rsa"
186        fake_report = mock.MagicMock()
187        instance_object = mock.MagicMock()
188        instance_object.ip = "1.1.1.1"
189        instance_object.vnc_port = 9999
190        instance_object.adb_port = "9999"
191        instance_object.islocal = False
192        instance_object.ssh_tunnel_is_connected = False
193        self.Patch(utils, "AutoConnect")
194        self.Patch(reconnect, "StartVnc")
195        self.Patch(reconnect, "_IsWebrtcEnable", return_value=False)
196        #test reconnect remote instance when avd_type as gce.
197        instance_object.avd_type = "gce"
198        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report)
199        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
200                                             rsa_key_file=ssh_private_key_path,
201                                             target_vnc_port=constants.GCE_VNC_PORT,
202                                             target_adb_port=constants.GCE_ADB_PORT,
203                                             ssh_user=constants.GCE_USER,
204                                             extra_args_ssh_tunnel=None)
205        reconnect.StartVnc.assert_called_once()
206
207        #test reconnect remote instance when avd_type as cuttlefish.
208        instance_object.avd_type = "cuttlefish"
209        reconnect.StartVnc.call_count = 0
210        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report)
211        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
212                                             rsa_key_file=ssh_private_key_path,
213                                             target_vnc_port=constants.CF_VNC_PORT,
214                                             target_adb_port=constants.CF_ADB_PORT,
215                                             ssh_user=constants.GCE_USER,
216                                             extra_args_ssh_tunnel=None)
217        reconnect.StartVnc.assert_called_once()
218
219    def testReconnectInstanceUnknownAvdType(self):
220        """Test reconnect instances of unknown avd type."""
221        ssh_private_key_path = "/fake/acloud_rsa"
222        fake_report = mock.MagicMock()
223        instance_object = mock.MagicMock()
224        instance_object.avd_type = "unknown"
225        self.assertRaises(errors.UnknownAvdType,
226                          reconnect.ReconnectInstance,
227                          ssh_private_key_path,
228                          instance_object,
229                          fake_report)
230
231    def testReconnectInstanceNoAvdType(self):
232        """Test reconnect instances with no avd type."""
233        ssh_private_key_path = "/fake/acloud_rsa"
234        fake_report = mock.MagicMock()
235        instance_object = mock.MagicMock()
236        self.assertRaises(errors.UnknownAvdType,
237                          reconnect.ReconnectInstance,
238                          ssh_private_key_path,
239                          instance_object,
240                          fake_report)
241
242    def testStartVnc(self):
243        """Test start Vnc."""
244        self.Patch(subprocess, "check_call", return_value=True)
245        self.Patch(utils, "IsCommandRunning", return_value=False)
246        self.Patch(utils, "LaunchVncClient")
247        vnc_port = 5555
248        display = ""
249        reconnect.StartVnc(vnc_port, display)
250        utils.LaunchVncClient.assert_called_with(5555)
251
252        display = "888x777 (99)"
253        utils.AutoConnect.call_count = 0
254        reconnect.StartVnc(vnc_port, display)
255        utils.LaunchVncClient.assert_called_with(5555, "888", "777")
256
257
258if __name__ == "__main__":
259    unittest.main()
260