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 20from unittest import mock 21 22from acloud import errors 23from acloud.internal import constants 24from acloud.internal.lib import auth 25from acloud.internal.lib import android_compute_client 26from acloud.internal.lib import cvd_runtime_config 27from acloud.internal.lib import driver_test_lib 28from acloud.internal.lib import gcompute_client 29from acloud.internal.lib import utils 30from acloud.internal.lib import ssh as ssh_object 31from acloud.internal.lib.adb_tools import AdbTools 32from acloud.list import list as list_instance 33from acloud.public import config 34from acloud.reconnect import reconnect 35 36 37ForwardedPorts = collections.namedtuple("ForwardedPorts", 38 [constants.VNC_PORT, 39 constants.ADB_PORT, 40 constants.FASTBOOT_PORT]) 41 42 43class ReconnectTest(driver_test_lib.BaseDriverTest): 44 """Test reconnect functions.""" 45 46 # pylint: disable=no-member, too-many-statements 47 def testReconnectInstance(self): 48 """Test Reconnect Instances.""" 49 ssh_private_key_path = "/fake/acloud_rsa" 50 fake_report = mock.MagicMock() 51 instance_object = mock.MagicMock() 52 instance_object.name = "fake_name" 53 instance_object.ip = "1.1.1.1" 54 instance_object.islocal = False 55 instance_object.adb_port = "8686" 56 instance_object.fastboot_port = "9686" 57 instance_object.avd_type = "cuttlefish" 58 self.Patch(subprocess, "check_call", return_value=True) 59 self.Patch(utils, "LaunchVncClient") 60 self.Patch(utils, "AutoConnect") 61 self.Patch(AdbTools, "IsAdbConnected", return_value=False) 62 self.Patch(AdbTools, "IsAdbConnectionAlive", return_value=False) 63 self.Patch(utils, "IsCommandRunning", return_value=False) 64 fake_device_dict = { 65 constants.IP: "1.1.1.1", 66 constants.INSTANCE_NAME: "fake_name", 67 constants.VNC_PORT: 6666, 68 constants.ADB_PORT: "8686", 69 constants.FASTBOOT_PORT: "9686", 70 constants.DEVICE_SERIAL: "127.0.0.1:8686" 71 } 72 73 # test ssh tunnel not connected, remote instance. 74 instance_object.vnc_port = 6666 75 instance_object.display = "" 76 utils.AutoConnect.call_count = 0 77 reconnect.ReconnectInstance( 78 ssh_private_key_path, instance_object, fake_report, autoconnect="vnc") 79 utils.AutoConnect.assert_not_called() 80 utils.LaunchVncClient.assert_called_with(6666) 81 fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict) 82 83 instance_object.display = "888x777 (99)" 84 utils.AutoConnect.call_count = 0 85 reconnect.ReconnectInstance( 86 ssh_private_key_path, instance_object, fake_report, autoconnect="vnc") 87 utils.AutoConnect.assert_not_called() 88 utils.LaunchVncClient.assert_called_with(6666, "888", "777") 89 fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict) 90 91 # test ssh tunnel connected , remote instance. 92 instance_object.ssh_tunnel_is_connected = False 93 instance_object.display = "" 94 utils.AutoConnect.call_count = 0 95 instance_object.vnc_port = 5555 96 extra_args_ssh_tunnel = None 97 self.Patch(utils, "AutoConnect", 98 return_value=ForwardedPorts(vnc_port=11111, adb_port=22222, fastboot_port=33333)) 99 reconnect.ReconnectInstance( 100 ssh_private_key_path, instance_object, fake_report, autoconnect="vnc") 101 utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip, 102 rsa_key_file=ssh_private_key_path, 103 target_vnc_port=constants.CF_VNC_PORT, 104 target_adb_port=constants.CF_ADB_PORT, 105 target_fastboot_port=constants.CF_FASTBOOT_PORT, 106 ssh_user=constants.GCE_USER, 107 extra_args_ssh_tunnel=extra_args_ssh_tunnel) 108 utils.LaunchVncClient.assert_called_with(11111) 109 fake_device_dict = { 110 constants.IP: "1.1.1.1", 111 constants.INSTANCE_NAME: "fake_name", 112 constants.VNC_PORT: 11111, 113 constants.ADB_PORT: 22222, 114 constants.FASTBOOT_PORT: 33333, 115 constants.DEVICE_SERIAL: "127.0.0.1:22222" 116 } 117 fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict) 118 119 instance_object.display = "999x777 (99)" 120 extra_args_ssh_tunnel = "fake_extra_args_ssh_tunnel" 121 utils.AutoConnect.call_count = 0 122 reconnect.ReconnectInstance( 123 ssh_private_key_path, instance_object, fake_report, 124 extra_args_ssh_tunnel=extra_args_ssh_tunnel, 125 autoconnect="vnc") 126 utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip, 127 rsa_key_file=ssh_private_key_path, 128 target_vnc_port=constants.CF_VNC_PORT, 129 target_adb_port=constants.CF_ADB_PORT, 130 target_fastboot_port=constants.CF_FASTBOOT_PORT, 131 ssh_user=constants.GCE_USER, 132 extra_args_ssh_tunnel=extra_args_ssh_tunnel) 133 utils.LaunchVncClient.assert_called_with(11111, "999", "777") 134 fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict) 135 136 # test fail reconnect report. 137 self.Patch(utils, "AutoConnect", 138 return_value=ForwardedPorts(vnc_port=None, adb_port=None, fastboot_port=None)) 139 reconnect.ReconnectInstance( 140 ssh_private_key_path, instance_object, fake_report, autoconnect="vnc") 141 fake_device_dict = { 142 constants.IP: "1.1.1.1", 143 constants.INSTANCE_NAME: "fake_name", 144 constants.VNC_PORT: None, 145 constants.ADB_PORT: None, 146 constants.FASTBOOT_PORT: None 147 } 148 fake_report.AddData.assert_called_with(key="device_failing_reconnect", 149 value=fake_device_dict) 150 151 # test reconnect local instance. 152 instance_object.islocal = True 153 instance_object.display = "" 154 instance_object.vnc_port = 5555 155 instance_object.ssh_tunnel_is_connected = False 156 utils.AutoConnect.call_count = 0 157 reconnect.ReconnectInstance( 158 ssh_private_key_path, instance_object, fake_report, autoconnect="vnc") 159 utils.AutoConnect.assert_not_called() 160 utils.LaunchVncClient.assert_called_with(5555) 161 fake_device_dict = { 162 constants.IP: "1.1.1.1", 163 constants.INSTANCE_NAME: "fake_name", 164 constants.VNC_PORT: 5555, 165 constants.ADB_PORT: "8686", 166 constants.FASTBOOT_PORT: "9686" 167 } 168 fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict) 169 170 # pylint: disable=no-member 171 def testReconnectInstanceWithWebRTC(self): 172 """Test reconnect instances with WebRTC.""" 173 ssh_private_key_path = "/fake/acloud_rsa" 174 fake_report = mock.MagicMock() 175 instance_object = mock.MagicMock() 176 instance_object.ip = "1.1.1.1" 177 instance_object.islocal = False 178 instance_object.adb_port = "8686" 179 instance_object.avd_type = "cuttlefish" 180 self.Patch(subprocess, "check_call", return_value=True) 181 self.Patch(utils, "LaunchVncClient") 182 self.Patch(utils, "AutoConnect") 183 self.Patch(utils, "LaunchBrowser") 184 self.Patch(utils, "GetWebrtcPortFromSSHTunnel", return_value=None) 185 self.Patch(utils, "EstablishWebRTCSshTunnel") 186 self.Patch(utils, "PickFreePort", return_value=12345) 187 self.Patch(AdbTools, "IsAdbConnected", return_value=False) 188 self.Patch(AdbTools, "IsAdbConnectionAlive", return_value=False) 189 self.Patch(utils, "IsCommandRunning", return_value=False) 190 191 # test ssh tunnel not reconnect to the remote instance. 192 instance_object.vnc_port = 6666 193 instance_object.display = "" 194 utils.AutoConnect.call_count = 0 195 reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report, 196 None, "webrtc") 197 utils.AutoConnect.assert_not_called() 198 utils.LaunchVncClient.assert_not_called() 199 utils.EstablishWebRTCSshTunnel.assert_called_with(extra_args_ssh_tunnel=None, 200 webrtc_local_port=12345, 201 ip_addr='1.1.1.1', 202 rsa_key_file='/fake/acloud_rsa', 203 ssh_user='vsoc-01') 204 utils.LaunchBrowser.assert_called_with('localhost', 12345) 205 utils.PickFreePort.assert_called_once() 206 utils.PickFreePort.reset_mock() 207 208 self.Patch(utils, "GetWebrtcPortFromSSHTunnel", return_value="11111") 209 reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report, 210 None, "webrtc") 211 utils.PickFreePort.assert_not_called() 212 213 # local webrtc instance 214 instance_object.islocal = True 215 reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report, 216 None, "webrtc") 217 utils.PickFreePort.assert_not_called() 218 219 # autoconnect adb only should launch nothing. 220 utils.LaunchBrowser.reset_mock() 221 utils.LaunchVncClient.reset_mock() 222 reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report, 223 None, "adb") 224 utils.LaunchBrowser.assert_not_called() 225 utils.LaunchVncClient.assert_not_called() 226 227 228 def testReconnectInstanceAvdtype(self): 229 """Test Reconnect Instances of avd_type.""" 230 ssh_private_key_path = "/fake/acloud_rsa" 231 fake_report = mock.MagicMock() 232 instance_object = mock.MagicMock() 233 instance_object.ip = "1.1.1.1" 234 instance_object.vnc_port = 9999 235 instance_object.adb_port = "9999" 236 instance_object.islocal = False 237 instance_object.ssh_tunnel_is_connected = False 238 self.Patch(utils, "AutoConnect") 239 self.Patch(reconnect, "StartVnc") 240 #test reconnect remote instance when avd_type as gce. 241 instance_object.avd_type = "gce" 242 reconnect.ReconnectInstance( 243 ssh_private_key_path, instance_object, fake_report, autoconnect="vnc") 244 utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip, 245 rsa_key_file=ssh_private_key_path, 246 target_vnc_port=constants.GCE_VNC_PORT, 247 target_adb_port=constants.GCE_ADB_PORT, 248 target_fastboot_port=None, 249 ssh_user=constants.GCE_USER, 250 extra_args_ssh_tunnel=None) 251 reconnect.StartVnc.assert_called_once() 252 253 #test reconnect remote instance when avd_type as cuttlefish. 254 instance_object.avd_type = "cuttlefish" 255 reconnect.StartVnc.call_count = 0 256 reconnect.ReconnectInstance( 257 ssh_private_key_path, instance_object, fake_report, autoconnect="vnc") 258 utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip, 259 rsa_key_file=ssh_private_key_path, 260 target_vnc_port=constants.CF_VNC_PORT, 261 target_adb_port=constants.CF_ADB_PORT, 262 target_fastboot_port=constants.CF_FASTBOOT_PORT, 263 ssh_user=constants.GCE_USER, 264 extra_args_ssh_tunnel=None) 265 reconnect.StartVnc.assert_called_once() 266 267 def testReconnectInstanceUnknownAvdType(self): 268 """Test reconnect instances of unknown avd type.""" 269 ssh_private_key_path = "/fake/acloud_rsa" 270 fake_report = mock.MagicMock() 271 instance_object = mock.MagicMock() 272 instance_object.avd_type = "unknown" 273 self.assertRaises(errors.UnknownAvdType, 274 reconnect.ReconnectInstance, 275 ssh_private_key_path, 276 instance_object, 277 fake_report) 278 279 def testReconnectInstanceNoAvdType(self): 280 """Test reconnect instances with no avd type.""" 281 ssh_private_key_path = "/fake/acloud_rsa" 282 fake_report = mock.MagicMock() 283 instance_object = mock.MagicMock() 284 self.assertRaises(errors.UnknownAvdType, 285 reconnect.ReconnectInstance, 286 ssh_private_key_path, 287 instance_object, 288 fake_report) 289 290 def testStartVnc(self): 291 """Test start Vnc.""" 292 self.Patch(subprocess, "check_call", return_value=True) 293 self.Patch(utils, "IsCommandRunning", return_value=False) 294 self.Patch(utils, "LaunchVncClient") 295 vnc_port = 5555 296 display = "" 297 reconnect.StartVnc(vnc_port, display) 298 utils.LaunchVncClient.assert_called_with(5555) 299 300 display = "888x777 (99)" 301 utils.AutoConnect.call_count = 0 302 reconnect.StartVnc(vnc_port, display) 303 utils.LaunchVncClient.assert_called_with(5555, "888", "777") 304 utils.LaunchVncClient.reset_mock() 305 306 self.Patch(utils, "IsCommandRunning", return_value=True) 307 reconnect.StartVnc(vnc_port, display) 308 utils.LaunchVncClient.assert_not_called() 309 310 # pylint: disable=protected-access 311 def testIsWebrtcEnable(self): 312 """Test _IsWebrtcEnable.""" 313 fake_ins = mock.MagicMock() 314 fake_ins.islocal = True 315 fake_ins.cf_runtime_cfg = mock.MagicMock() 316 fake_ins.cf_runtime_cfg.enable_webrtc = False 317 reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", "") 318 self.assertFalse(reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", "")) 319 320 fake_ins.islocal = False 321 fake_runtime_config = mock.MagicMock() 322 fake_runtime_config.enable_webrtc = True 323 self.Patch(ssh_object, "Ssh") 324 self.Patch(ssh_object.Ssh, "GetCmdOutput", return_value="fake_rawdata") 325 self.Patch(cvd_runtime_config, "CvdRuntimeConfig", 326 return_value=fake_runtime_config) 327 self.assertTrue(reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", "")) 328 329 self.Patch(cvd_runtime_config, "CvdRuntimeConfig", 330 side_effect=errors.ConfigError) 331 self.assertFalse(reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", "")) 332 333 def testRun(self): 334 """Test Run.""" 335 fake_args = mock.MagicMock() 336 fake_args.autoconnect = "webrtc" 337 fake_args.instance_names = ["fake-ins-name"] 338 fake_ins1 = mock.MagicMock() 339 fake_ins1.avd_type = "cuttlefish" 340 fake_ins1.islocal = False 341 fake_ins2 = mock.MagicMock() 342 fake_ins2.avd_type = "cuttlefish" 343 fake_ins2.islocal = False 344 fake_ins_gf = mock.MagicMock() 345 fake_ins_gf.avd_type = "goldfish" 346 fake_ins_gf.islocal = False 347 fake_ins_gf.vnc_port = 1234 348 ins_to_reconnect = [fake_ins1] 349 # mock args.all equal to True and return 3 instances. 350 all_ins_to_reconnect = [fake_ins1, fake_ins2, fake_ins_gf] 351 cfg = mock.MagicMock() 352 cfg.ssh_private_key_path = None 353 cfg.extra_args_ssh_tunnel = None 354 self.Patch(config, "GetAcloudConfig", return_value=cfg) 355 self.Patch(list_instance, "GetInstancesFromInstanceNames", 356 return_value=ins_to_reconnect) 357 self.Patch(list_instance, "ChooseInstances", 358 return_value=all_ins_to_reconnect) 359 self.Patch(auth, "CreateCredentials") 360 self.Patch(android_compute_client, "AndroidComputeClient") 361 self.Patch(android_compute_client.AndroidComputeClient, 362 "AddSshRsaInstanceMetadata") 363 self.Patch(reconnect, "ReconnectInstance") 364 365 reconnect.Run(fake_args) 366 list_instance.GetInstancesFromInstanceNames.assert_called_once() 367 self.assertEqual(reconnect.ReconnectInstance.call_count, 1) 368 reconnect.ReconnectInstance.reset_mock() 369 370 # should reconnect all instances 371 fake_args.instance_names = None 372 reconnect.Run(fake_args) 373 list_instance.ChooseInstances.assert_called_once() 374 self.assertEqual(reconnect.ReconnectInstance.call_count, 3) 375 reconnect.ReconnectInstance.reset_mock() 376 377 fake_ins1.islocal = True 378 fake_ins2.avd_type = "unknown" 379 self.Patch(list_instance, "ChooseInstances", 380 return_value=[fake_ins1, fake_ins2]) 381 reconnect.Run(fake_args) 382 self.assertEqual(reconnect.ReconnectInstance.call_count, 1) 383 384 def testGetSshConnectHostname(self): 385 """Test GetSshConnectHostname.""" 386 self.Patch(gcompute_client, "GetGCEHostName", return_value="fake_host") 387 instance = mock.MagicMock() 388 instance.islocal = True 389 cfg = mock.MagicMock() 390 self.assertEqual(None, reconnect.GetSshConnectHostname(cfg, instance)) 391 392 # Remote instance will get the GCE hostname. 393 instance.islocal = False 394 cfg.connect_hostname = True 395 self.assertEqual("fake_host", 396 reconnect.GetSshConnectHostname(cfg, instance)) 397 398 399if __name__ == "__main__": 400 unittest.main() 401