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