xref: /aosp_15_r20/tools/acloud/reconnect/reconnect_test.py (revision 800a58d989c669b8eb8a71d8df53b1ba3d411444)

# Copyright 2018 - The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for reconnect."""

import collections
import unittest
import subprocess

from unittest import mock

from acloud import errors
from acloud.internal import constants
from acloud.internal.lib import auth
from acloud.internal.lib import android_compute_client
from acloud.internal.lib import cvd_runtime_config
from acloud.internal.lib import driver_test_lib
from acloud.internal.lib import gcompute_client
from acloud.internal.lib import utils
from acloud.internal.lib import ssh as ssh_object
from acloud.internal.lib.adb_tools import AdbTools
from acloud.list import list as list_instance
from acloud.public import config
from acloud.reconnect import reconnect


ForwardedPorts = collections.namedtuple("ForwardedPorts",
                                        [constants.VNC_PORT, constants.ADB_PORT])


class ReconnectTest(driver_test_lib.BaseDriverTest):
    """Test reconnect functions."""

    # pylint: disable=no-member, too-many-statements
    def testReconnectInstance(self):
        """Test Reconnect Instances."""
        ssh_private_key_path = "/fake/acloud_rsa"
        fake_report = mock.MagicMock()
        instance_object = mock.MagicMock()
        instance_object.name = "fake_name"
        instance_object.ip = "1.1.1.1"
        instance_object.islocal = False
        instance_object.adb_port = "8686"
        instance_object.avd_type = "cuttlefish"
        self.Patch(subprocess, "check_call", return_value=True)
        self.Patch(utils, "LaunchVncClient")
        self.Patch(utils, "AutoConnect")
        self.Patch(AdbTools, "IsAdbConnected", return_value=False)
        self.Patch(AdbTools, "IsAdbConnectionAlive", return_value=False)
        self.Patch(utils, "IsCommandRunning", return_value=False)
        fake_device_dict = {
            constants.IP: "1.1.1.1",
            constants.INSTANCE_NAME: "fake_name",
            constants.VNC_PORT: 6666,
            constants.ADB_PORT: "8686",
            constants.DEVICE_SERIAL: "127.0.0.1:8686"
        }

        # test ssh tunnel not connected, remote instance.
        instance_object.vnc_port = 6666
        instance_object.display = ""
        utils.AutoConnect.call_count = 0
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report, autoconnect="vnc")
        utils.AutoConnect.assert_not_called()
        utils.LaunchVncClient.assert_called_with(6666)
        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)

        instance_object.display = "888x777 (99)"
        utils.AutoConnect.call_count = 0
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report, autoconnect="vnc")
        utils.AutoConnect.assert_not_called()
        utils.LaunchVncClient.assert_called_with(6666, "888", "777")
        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)

        # test ssh tunnel connected , remote instance.
        instance_object.ssh_tunnel_is_connected = False
        instance_object.display = ""
        utils.AutoConnect.call_count = 0
        instance_object.vnc_port = 5555
        extra_args_ssh_tunnel = None
        self.Patch(utils, "AutoConnect",
                   return_value=ForwardedPorts(vnc_port=11111, adb_port=22222))
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report, autoconnect="vnc")
        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
                                             rsa_key_file=ssh_private_key_path,
                                             target_vnc_port=constants.CF_VNC_PORT,
                                             target_adb_port=constants.CF_ADB_PORT,
                                             ssh_user=constants.GCE_USER,
                                             extra_args_ssh_tunnel=extra_args_ssh_tunnel)
        utils.LaunchVncClient.assert_called_with(11111)
        fake_device_dict = {
            constants.IP: "1.1.1.1",
            constants.INSTANCE_NAME: "fake_name",
            constants.VNC_PORT: 11111,
            constants.ADB_PORT: 22222,
            constants.DEVICE_SERIAL: "127.0.0.1:22222"
        }
        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)

        instance_object.display = "999x777 (99)"
        extra_args_ssh_tunnel = "fake_extra_args_ssh_tunnel"
        utils.AutoConnect.call_count = 0
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report,
            extra_args_ssh_tunnel=extra_args_ssh_tunnel,
            autoconnect="vnc")
        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
                                             rsa_key_file=ssh_private_key_path,
                                             target_vnc_port=constants.CF_VNC_PORT,
                                             target_adb_port=constants.CF_ADB_PORT,
                                             ssh_user=constants.GCE_USER,
                                             extra_args_ssh_tunnel=extra_args_ssh_tunnel)
        utils.LaunchVncClient.assert_called_with(11111, "999", "777")
        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)

        # test fail reconnect report.
        self.Patch(utils, "AutoConnect",
                   return_value=ForwardedPorts(vnc_port=None, adb_port=None))
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report, autoconnect="vnc")
        fake_device_dict = {
            constants.IP: "1.1.1.1",
            constants.INSTANCE_NAME: "fake_name",
            constants.VNC_PORT: None,
            constants.ADB_PORT: None
        }
        fake_report.AddData.assert_called_with(key="device_failing_reconnect",
                                               value=fake_device_dict)

        # test reconnect local instance.
        instance_object.islocal = True
        instance_object.display = ""
        instance_object.vnc_port = 5555
        instance_object.ssh_tunnel_is_connected = False
        utils.AutoConnect.call_count = 0
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report, autoconnect="vnc")
        utils.AutoConnect.assert_not_called()
        utils.LaunchVncClient.assert_called_with(5555)
        fake_device_dict = {
            constants.IP: "1.1.1.1",
            constants.INSTANCE_NAME: "fake_name",
            constants.VNC_PORT: 5555,
            constants.ADB_PORT: "8686"
        }
        fake_report.AddData.assert_called_with(key="devices", value=fake_device_dict)

    # pylint: disable=no-member
    def testReconnectInstanceWithWebRTC(self):
        """Test reconnect instances with WebRTC."""
        ssh_private_key_path = "/fake/acloud_rsa"
        fake_report = mock.MagicMock()
        instance_object = mock.MagicMock()
        instance_object.ip = "1.1.1.1"
        instance_object.islocal = False
        instance_object.adb_port = "8686"
        instance_object.avd_type = "cuttlefish"
        self.Patch(subprocess, "check_call", return_value=True)
        self.Patch(utils, "LaunchVncClient")
        self.Patch(utils, "AutoConnect")
        self.Patch(utils, "LaunchBrowser")
        self.Patch(utils, "GetWebrtcPortFromSSHTunnel", return_value=None)
        self.Patch(utils, "EstablishWebRTCSshTunnel")
        self.Patch(utils, "PickFreePort", return_value=12345)
        self.Patch(AdbTools, "IsAdbConnected", return_value=False)
        self.Patch(AdbTools, "IsAdbConnectionAlive", return_value=False)
        self.Patch(utils, "IsCommandRunning", return_value=False)

        # test ssh tunnel not reconnect to the remote instance.
        instance_object.vnc_port = 6666
        instance_object.display = ""
        utils.AutoConnect.call_count = 0
        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report,
                                    None, "webrtc")
        utils.AutoConnect.assert_not_called()
        utils.LaunchVncClient.assert_not_called()
        utils.EstablishWebRTCSshTunnel.assert_called_with(extra_args_ssh_tunnel=None,
                                                          webrtc_local_port=12345,
                                                          ip_addr='1.1.1.1',
                                                          rsa_key_file='/fake/acloud_rsa',
                                                          ssh_user='vsoc-01')
        utils.LaunchBrowser.assert_called_with('localhost', 12345)
        utils.PickFreePort.assert_called_once()
        utils.PickFreePort.reset_mock()

        self.Patch(utils, "GetWebrtcPortFromSSHTunnel", return_value="11111")
        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report,
                                    None, "webrtc")
        utils.PickFreePort.assert_not_called()

        # local webrtc instance
        instance_object.islocal = True
        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report,
                                    None, "webrtc")
        utils.PickFreePort.assert_not_called()

        # autoconnect adb only should launch nothing.
        utils.LaunchBrowser.reset_mock()
        utils.LaunchVncClient.reset_mock()
        reconnect.ReconnectInstance(ssh_private_key_path, instance_object, fake_report,
                                    None, "adb")
        utils.LaunchBrowser.assert_not_called()
        utils.LaunchVncClient.assert_not_called()


    def testReconnectInstanceAvdtype(self):
        """Test Reconnect Instances of avd_type."""
        ssh_private_key_path = "/fake/acloud_rsa"
        fake_report = mock.MagicMock()
        instance_object = mock.MagicMock()
        instance_object.ip = "1.1.1.1"
        instance_object.vnc_port = 9999
        instance_object.adb_port = "9999"
        instance_object.islocal = False
        instance_object.ssh_tunnel_is_connected = False
        self.Patch(utils, "AutoConnect")
        self.Patch(reconnect, "StartVnc")
        #test reconnect remote instance when avd_type as gce.
        instance_object.avd_type = "gce"
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report, autoconnect="vnc")
        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
                                             rsa_key_file=ssh_private_key_path,
                                             target_vnc_port=constants.GCE_VNC_PORT,
                                             target_adb_port=constants.GCE_ADB_PORT,
                                             ssh_user=constants.GCE_USER,
                                             extra_args_ssh_tunnel=None)
        reconnect.StartVnc.assert_called_once()

        #test reconnect remote instance when avd_type as cuttlefish.
        instance_object.avd_type = "cuttlefish"
        reconnect.StartVnc.call_count = 0
        reconnect.ReconnectInstance(
            ssh_private_key_path, instance_object, fake_report, autoconnect="vnc")
        utils.AutoConnect.assert_called_with(ip_addr=instance_object.ip,
                                             rsa_key_file=ssh_private_key_path,
                                             target_vnc_port=constants.CF_VNC_PORT,
                                             target_adb_port=constants.CF_ADB_PORT,
                                             ssh_user=constants.GCE_USER,
                                             extra_args_ssh_tunnel=None)
        reconnect.StartVnc.assert_called_once()

    def testReconnectInstanceUnknownAvdType(self):
        """Test reconnect instances of unknown avd type."""
        ssh_private_key_path = "/fake/acloud_rsa"
        fake_report = mock.MagicMock()
        instance_object = mock.MagicMock()
        instance_object.avd_type = "unknown"
        self.assertRaises(errors.UnknownAvdType,
                          reconnect.ReconnectInstance,
                          ssh_private_key_path,
                          instance_object,
                          fake_report)

    def testReconnectInstanceNoAvdType(self):
        """Test reconnect instances with no avd type."""
        ssh_private_key_path = "/fake/acloud_rsa"
        fake_report = mock.MagicMock()
        instance_object = mock.MagicMock()
        self.assertRaises(errors.UnknownAvdType,
                          reconnect.ReconnectInstance,
                          ssh_private_key_path,
                          instance_object,
                          fake_report)

    def testStartVnc(self):
        """Test start Vnc."""
        self.Patch(subprocess, "check_call", return_value=True)
        self.Patch(utils, "IsCommandRunning", return_value=False)
        self.Patch(utils, "LaunchVncClient")
        vnc_port = 5555
        display = ""
        reconnect.StartVnc(vnc_port, display)
        utils.LaunchVncClient.assert_called_with(5555)

        display = "888x777 (99)"
        utils.AutoConnect.call_count = 0
        reconnect.StartVnc(vnc_port, display)
        utils.LaunchVncClient.assert_called_with(5555, "888", "777")
        utils.LaunchVncClient.reset_mock()

        self.Patch(utils, "IsCommandRunning", return_value=True)
        reconnect.StartVnc(vnc_port, display)
        utils.LaunchVncClient.assert_not_called()

    # pylint: disable=protected-access
    def testIsWebrtcEnable(self):
        """Test _IsWebrtcEnable."""
        fake_ins = mock.MagicMock()
        fake_ins.islocal = True
        fake_ins.cf_runtime_cfg = mock.MagicMock()
        fake_ins.cf_runtime_cfg.enable_webrtc = False
        reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", "")
        self.assertFalse(reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", ""))

        fake_ins.islocal = False
        fake_runtime_config = mock.MagicMock()
        fake_runtime_config.enable_webrtc = True
        self.Patch(ssh_object, "Ssh")
        self.Patch(ssh_object.Ssh, "GetCmdOutput", return_value="fake_rawdata")
        self.Patch(cvd_runtime_config, "CvdRuntimeConfig",
                   return_value=fake_runtime_config)
        self.assertTrue(reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", ""))

        self.Patch(cvd_runtime_config, "CvdRuntimeConfig",
                   side_effect=errors.ConfigError)
        self.assertFalse(reconnect._IsWebrtcEnable(fake_ins, "fake_user", "ssh_pkey_path", ""))

    def testRun(self):
        """Test Run."""
        fake_args = mock.MagicMock()
        fake_args.autoconnect = "webrtc"
        fake_args.instance_names = ["fake-ins-name"]
        fake_ins1 = mock.MagicMock()
        fake_ins1.avd_type = "cuttlefish"
        fake_ins1.islocal = False
        fake_ins2 = mock.MagicMock()
        fake_ins2.avd_type = "cuttlefish"
        fake_ins2.islocal = False
        fake_ins_gf = mock.MagicMock()
        fake_ins_gf.avd_type = "goldfish"
        fake_ins_gf.islocal = False
        fake_ins_gf.vnc_port = 1234
        ins_to_reconnect = [fake_ins1]
        # mock args.all equal to True and return 3 instances.
        all_ins_to_reconnect = [fake_ins1, fake_ins2, fake_ins_gf]
        cfg = mock.MagicMock()
        cfg.ssh_private_key_path = None
        cfg.extra_args_ssh_tunnel = None
        self.Patch(config, "GetAcloudConfig", return_value=cfg)
        self.Patch(list_instance, "GetInstancesFromInstanceNames",
                   return_value=ins_to_reconnect)
        self.Patch(list_instance, "ChooseInstances",
                   return_value=all_ins_to_reconnect)
        self.Patch(auth, "CreateCredentials")
        self.Patch(android_compute_client, "AndroidComputeClient")
        self.Patch(android_compute_client.AndroidComputeClient,
                   "AddSshRsaInstanceMetadata")
        self.Patch(reconnect, "ReconnectInstance")

        reconnect.Run(fake_args)
        list_instance.GetInstancesFromInstanceNames.assert_called_once()
        self.assertEqual(reconnect.ReconnectInstance.call_count, 1)
        reconnect.ReconnectInstance.reset_mock()

        # should reconnect all instances
        fake_args.instance_names = None
        reconnect.Run(fake_args)
        list_instance.ChooseInstances.assert_called_once()
        self.assertEqual(reconnect.ReconnectInstance.call_count, 3)
        reconnect.ReconnectInstance.reset_mock()

        fake_ins1.islocal = True
        fake_ins2.avd_type = "unknown"
        self.Patch(list_instance, "ChooseInstances",
                   return_value=[fake_ins1, fake_ins2])
        reconnect.Run(fake_args)
        self.assertEqual(reconnect.ReconnectInstance.call_count, 1)

    def testGetSshConnectHostname(self):
        """Test GetSshConnectHostname."""
        self.Patch(gcompute_client, "GetGCEHostName", return_value="fake_host")
        instance = mock.MagicMock()
        instance.islocal = True
        cfg = mock.MagicMock()
        self.assertEqual(None, reconnect.GetSshConnectHostname(cfg, instance))

        # Remote instance will get the GCE hostname.
        instance.islocal = False
        cfg.connect_hostname = True
        self.assertEqual("fake_host",
                         reconnect.GetSshConnectHostname(cfg, instance))


if __name__ == "__main__":
    unittest.main()