xref: /aosp_15_r20/tools/acloud/restart/restart_test.py (revision 800a58d989c669b8eb8a71d8df53b1ba3d411444) 
1# Copyright 2021 - 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 restart."""
15import sys
16
17import unittest
18
19from unittest import mock
20
21from acloud import errors
22from acloud.internal.lib import driver_test_lib
23from acloud.internal.lib.ssh import Ssh
24from acloud.list import list as list_instances
25from acloud.powerwash import powerwash
26from acloud.public import config
27from acloud.public import report
28from acloud.reconnect import reconnect
29from acloud.restart import restart
30
31
32class RestartTest(driver_test_lib.BaseDriverTest):
33    """Test restart."""
34    @mock.patch.object(sys, "exit")
35    @mock.patch.object(restart, "RestartFromInstance")
36    def testRun(self, mock_restart, mock_exit):
37        """test Run."""
38        cfg = mock.MagicMock()
39        args = mock.MagicMock()
40        instance_obj = mock.MagicMock()
41        # Test case with provided instance name.
42        args.instance_name = "instance_1"
43        args.instance_id = 1
44        args.powerwash = False
45        self.Patch(config, "GetAcloudConfig", return_value=cfg)
46        self.Patch(list_instances, "GetInstancesFromInstanceNames",
47                   return_value=[instance_obj])
48        restart.Run(args)
49        mock_restart.assert_has_calls([
50            mock.call(cfg, instance_obj, args.instance_id, args.powerwash)])
51
52        # Test case for user select one instance to restart AVD.
53        selected_instance = mock.MagicMock()
54        self.Patch(list_instances, "GetCFRemoteInstances",
55                   return_value=selected_instance)
56        self.Patch(list_instances, "ChooseOneRemoteInstance",
57                   return_value=selected_instance)
58        args.instance_name = None
59        restart.Run(args)
60        mock_restart.assert_has_calls([
61            mock.call(cfg, selected_instance, args.instance_id, args.powerwash)])
62
63        # Test case for not support local instances.
64        local_instances = mock.MagicMock()
65        self.Patch(list_instances, "GetCFRemoteInstances",
66                   return_value=None)
67        self.Patch(list_instances, "GetLocalInstances",
68                   return_value=local_instances)
69        restart.Run(args)
70        mock_exit.assert_called_once()
71
72    # pylint: disable=no-member
73    def testRestartFromInstance(self):
74        """test RestartFromInstance."""
75        cfg = mock.MagicMock()
76        cfg.ssh_private_key_path = "fake_path"
77        cfg.extra_args_ssh_tunnel = ""
78        instance = mock.MagicMock()
79        instance.ip = "0.0.0.0"
80        instance.name = "ins-name"
81        self.Patch(powerwash, "PowerwashDevice")
82        self.Patch(reconnect, "ReconnectInstance")
83        self.Patch(report, "Report")
84        self.Patch(Ssh, "Run")
85        # should powerwash
86        restart.RestartFromInstance(cfg, instance, 1, True)
87        powerwash.PowerwashDevice.assert_called_once()
88
89        # should restart
90        restart.RestartFromInstance(cfg, instance, 1, False)
91        Ssh.Run.assert_called_once()
92
93        # coverage for except
94        self.Patch(Ssh, "Run",
95                   side_effect=errors.DeviceConnectionError())
96        restart.RestartFromInstance(cfg, instance, 1, False)
97
98
99if __name__ == '__main__':
100    unittest.main()
101