1#!/usr/bin/env vpython3 2# Copyright 2022 The Chromium Authors 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5"""Tests scenarios for ermine_ctl""" 6import logging 7import subprocess 8import time 9import unittest 10import unittest.mock as mock 11 12from base_ermine_ctl import BaseErmineCtl 13 14 15class BaseBaseErmineCtlTest(unittest.TestCase): 16 """Unit tests for BaseBaseErmineCtl interface.""" 17 18 def __init__(self, *args, **kwargs): 19 super().__init__(*args, **kwargs) 20 self.ermine_ctl = BaseErmineCtl() 21 22 def _set_mock_proc(self, return_value: int): 23 """Set |execute_command_async|'s return value to a mocked subprocess.""" 24 self.ermine_ctl.execute_command_async = mock.MagicMock() 25 mock_proc = mock.create_autospec(subprocess.Popen, instance=True) 26 mock_proc.communicate.return_value = 'foo', 'stderr' 27 mock_proc.returncode = return_value 28 self.ermine_ctl.execute_command_async.return_value = mock_proc 29 30 return mock_proc 31 32 def test_check_exists(self): 33 """Test |exists| returns True if tool command succeeds (returns 0).""" 34 self._set_mock_proc(return_value=0) 35 36 self.assertTrue(self.ermine_ctl.exists) 37 38 # Modifying this will not result in a change in state due to caching. 39 self._set_mock_proc(return_value=42) 40 self.assertTrue(self.ermine_ctl.exists) 41 42 def test_does_not_exist(self): 43 """Test |exists| returns False if tool command fails (returns != 0).""" 44 self._set_mock_proc(return_value=42) 45 46 self.assertFalse(self.ermine_ctl.exists) 47 48 def test_ready_raises_assertion_error_if_not_exist(self): 49 """Test |ready| raises AssertionError if tool does not exist.""" 50 self._set_mock_proc(return_value=42) 51 self.assertRaises(AssertionError, getattr, self.ermine_ctl, 'ready') 52 53 def test_ready_returns_false_if_bad_status(self): 54 """Test |ready| return False if tool has a bad status.""" 55 with mock.patch.object( 56 BaseErmineCtl, 'status', 57 new_callable=mock.PropertyMock) as mock_status, \ 58 mock.patch.object(BaseErmineCtl, 'exists', 59 new_callable=mock.PropertyMock) as mock_exists: 60 mock_exists.return_value = True 61 mock_status.return_value = (1, 'FakeStatus') 62 self.assertFalse(self.ermine_ctl.ready) 63 64 def test_ready_returns_true(self): 65 """Test |ready| return True if tool returns good status (rc = 0).""" 66 with mock.patch.object( 67 BaseErmineCtl, 'status', 68 new_callable=mock.PropertyMock) as mock_status, \ 69 mock.patch.object(BaseErmineCtl, 'exists', 70 new_callable=mock.PropertyMock) as mock_exists: 71 mock_exists.return_value = True 72 mock_status.return_value = (0, 'FakeStatus') 73 self.assertTrue(self.ermine_ctl.ready) 74 75 def test_status_raises_assertion_error_if_dne(self): 76 """Test |status| returns |InvalidState| if tool does not exist.""" 77 with mock.patch.object(BaseErmineCtl, 78 'exists', 79 new_callable=mock.PropertyMock) as mock_exists: 80 mock_exists.return_value = False 81 82 self.assertRaises(AssertionError, getattr, self.ermine_ctl, 83 'status') 84 85 def test_status_returns_rc_and_stdout(self): 86 """Test |status| returns subprocess stdout and rc if tool exists.""" 87 with mock.patch.object(BaseErmineCtl, 88 'exists', 89 new_callable=mock.PropertyMock) as _: 90 self._set_mock_proc(return_value=10) 91 92 self.assertEqual(self.ermine_ctl.status, (10, 'foo')) 93 94 def test_status_returns_timeout_state(self): 95 """Test |status| returns |Timeout| if exception is raised.""" 96 with mock.patch.object( 97 BaseErmineCtl, 'exists', new_callable=mock.PropertyMock) as _, \ 98 mock.patch.object(logging, 'warning') as _: 99 mock_proc = self._set_mock_proc(return_value=0) 100 mock_proc.wait.side_effect = subprocess.TimeoutExpired( 101 'cmd', 'some timeout') 102 103 self.assertEqual(self.ermine_ctl.status, (-1, 'Timeout')) 104 105 def test_wait_until_ready_raises_assertion_error_if_tool_dne(self): 106 """Test |wait_until_ready| is returns false if tool does not exist.""" 107 with mock.patch.object(BaseErmineCtl, 108 'exists', 109 new_callable=mock.PropertyMock) as mock_exists: 110 mock_exists.return_value = False 111 112 self.assertRaises(AssertionError, self.ermine_ctl.wait_until_ready) 113 114 def test_wait_until_ready_loops_until_ready(self): 115 """Test |wait_until_ready| loops until |ready| returns True.""" 116 with mock.patch.object(BaseErmineCtl, 'exists', 117 new_callable=mock.PropertyMock) as mock_exists, \ 118 mock.patch.object(time, 'sleep') as mock_sleep, \ 119 mock.patch.object(BaseErmineCtl, 'ready', 120 new_callable=mock.PropertyMock) as mock_ready: 121 mock_exists.return_value = True 122 mock_ready.side_effect = [False, False, False, True] 123 124 self.ermine_ctl.wait_until_ready() 125 126 self.assertEqual(mock_ready.call_count, 4) 127 self.assertEqual(mock_sleep.call_count, 3) 128 129 def test_wait_until_ready_raises_assertion_error_if_attempts_exceeded( 130 self): 131 """Test |wait_until_ready| loops if |ready| is not True n attempts.""" 132 with mock.patch.object(BaseErmineCtl, 'exists', 133 new_callable=mock.PropertyMock) as mock_exists, \ 134 mock.patch.object(time, 'sleep') as mock_sleep, \ 135 mock.patch.object(BaseErmineCtl, 'ready', 136 new_callable=mock.PropertyMock) as mock_ready: 137 mock_exists.return_value = True 138 mock_ready.side_effect = [False] * 15 + [True] 139 140 self.assertRaises(TimeoutError, self.ermine_ctl.wait_until_ready) 141 142 self.assertEqual(mock_ready.call_count, 10) 143 self.assertEqual(mock_sleep.call_count, 10) 144 145 def test_take_to_shell_raises_assertion_error_if_tool_dne(self): 146 """Test |take_to_shell| throws AssertionError if not ready is False.""" 147 with mock.patch.object(BaseErmineCtl, 148 'exists', 149 new_callable=mock.PropertyMock) as mock_exists: 150 mock_exists.return_value = False 151 self.assertRaises(AssertionError, self.ermine_ctl.take_to_shell) 152 153 def test_take_to_shell_exits_on_complete_state(self): 154 """Test |take_to_shell| exits with no calls if in completed state.""" 155 with mock.patch.object(BaseErmineCtl, 156 'wait_until_ready') as mock_wait_ready, \ 157 mock.patch.object( 158 BaseErmineCtl, 'status', 159 new_callable=mock.PropertyMock) as mock_status: 160 mock_proc = self._set_mock_proc(return_value=52) 161 mock_wait_ready.return_value = True 162 mock_status.return_value = (0, 'Shell') 163 164 self.ermine_ctl.take_to_shell() 165 166 self.assertEqual(mock_proc.call_count, 0) 167 168 def test_take_to_shell_invalid_state_raises_not_implemented_error(self): 169 """Test |take_to_shell| raises exception if invalid state is returned. 170 """ 171 with mock.patch.object(BaseErmineCtl, 172 'wait_until_ready') as mock_wait_ready, \ 173 mock.patch.object( 174 BaseErmineCtl, 'status', 175 new_callable=mock.PropertyMock) as mock_status: 176 mock_wait_ready.return_value = True 177 mock_status.return_value = (0, 'SomeUnknownState') 178 179 self.assertRaises(NotImplementedError, 180 self.ermine_ctl.take_to_shell) 181 182 def test_take_to_shell_with_max_transitions_raises_runtime_error(self): 183 """Test |take_to_shell| raises exception on too many transitions. 184 185 |take_to_shell| attempts to transition from one state to another. 186 After 5 attempts, if this does not end in the completed state, an 187 Exception is thrown. 188 """ 189 with mock.patch.object(BaseErmineCtl, 190 'wait_until_ready') as mock_wait_ready, \ 191 mock.patch.object( 192 BaseErmineCtl, 'status', 193 new_callable=mock.PropertyMock) as mock_status: 194 mock_wait_ready.return_value = True 195 # Returns too many state transitions before CompleteState. 196 mock_status.side_effect = [(0, 'Unknown'), 197 (0, 'KnownWithPassword'), 198 (0, 'Unknown')] * 3 + [ 199 (0, 'CompleteState') 200 ] 201 self.assertRaises(RuntimeError, self.ermine_ctl.take_to_shell) 202 203 def test_take_to_shell_executes_known_commands(self): 204 """Test |take_to_shell| executes commands if necessary. 205 206 Some states can only be transitioned between with specific commands. 207 These are executed by |take_to_shell| until the final test |Shell| is 208 reached. 209 """ 210 with mock.patch.object(BaseErmineCtl, 211 'wait_until_ready') as mock_wait_ready, \ 212 mock.patch.object( 213 BaseErmineCtl, 'status', 214 new_callable=mock.PropertyMock) as mock_status: 215 self._set_mock_proc(return_value=0) 216 mock_wait_ready.return_value = True 217 mock_status.side_effect = [(0, 'Unknown'), (0, 'SetPassword'), 218 (0, 'Shell')] 219 220 self.ermine_ctl.take_to_shell() 221 222 self.assertEqual(self.ermine_ctl.execute_command_async.call_count, 223 2) 224 self.ermine_ctl.execute_command_async.assert_has_calls([ 225 mock.call(['erminectl', 'oobe', 'skip']), 226 mock.call().communicate(), 227 mock.call([ 228 'erminectl', 'oobe', 'set_password', 229 'workstation_test_password' 230 ]), 231 mock.call().communicate() 232 ]) 233 234 235if __name__ == '__main__': 236 unittest.main() 237