ermine_ctl_test.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255
  1. #!/usr/bin/env vpython3
  2. # Copyright 2022 The Chromium Authors. All rights reserved.
  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"""
  6. import common
  7. import os
  8. import logging
  9. import subprocess
  10. import time
  11. import unittest
  12. import unittest.mock as mock
  13. from target import Target
  14. from ermine_ctl import ErmineCtl
  15. class TestDiscoverDeviceTarget(unittest.TestCase):
  16. def setUp(self):
  17. self.mock_target = mock.create_autospec(Target, instance=True)
  18. self.ermine_ctl = ErmineCtl(self.mock_target)
  19. def testCheckExists(self):
  20. self.mock_target.RunCommand.return_value = 0
  21. self.assertFalse(self.ermine_ctl._ermine_exists)
  22. self.assertFalse(self.ermine_ctl._ermine_exists_check)
  23. self.assertTrue(self.ermine_ctl.exists)
  24. self.assertTrue(self.ermine_ctl._ermine_exists)
  25. self.assertTrue(self.ermine_ctl._ermine_exists_check)
  26. # Modifying this will not result in a change in state.
  27. self.mock_target.RunCommand.return_value = 42
  28. self.assertTrue(self.ermine_ctl.exists)
  29. self.assertTrue(self.ermine_ctl._ermine_exists)
  30. self.assertTrue(self.ermine_ctl._ermine_exists_check)
  31. def testDoesNotExist(self):
  32. self.mock_target.RunCommand.return_value = 42
  33. self.assertFalse(self.ermine_ctl._ermine_exists)
  34. self.assertFalse(self.ermine_ctl._ermine_exists_check)
  35. self.assertFalse(self.ermine_ctl.exists)
  36. self.assertFalse(self.ermine_ctl._ermine_exists)
  37. self.assertTrue(self.ermine_ctl._ermine_exists_check)
  38. def testReadyReturnsFalseIfNoExist(self):
  39. self.ermine_ctl._ermine_exists = False
  40. self.ermine_ctl._ermine_exists_check = True
  41. self.assertFalse(self.ermine_ctl.ready)
  42. def testReadyReturnsFalseIfBadStatus(self):
  43. with mock.patch.object(
  44. ErmineCtl, 'status', new_callable=mock.PropertyMock) as mock_status, \
  45. mock.patch.object(ErmineCtl, 'exists',
  46. new_callable=mock.PropertyMock) as mock_exists:
  47. mock_exists.return_value = True
  48. mock_status.return_value = (1, 'FakeStatus')
  49. self.assertFalse(self.ermine_ctl.ready)
  50. def testReadyReturnsTrue(self):
  51. with mock.patch.object(
  52. ErmineCtl, 'status', new_callable=mock.PropertyMock) as mock_status, \
  53. mock.patch.object(ErmineCtl, 'exists',
  54. new_callable=mock.PropertyMock) as mock_exists:
  55. mock_exists.return_value = True
  56. mock_status.return_value = (0, 'FakeStatus')
  57. self.assertTrue(self.ermine_ctl.ready)
  58. def testStatusReturnsInvalidStateIfNoExists(self):
  59. with mock.patch.object(ErmineCtl, 'exists',
  60. new_callable=mock.PropertyMock) as mock_exists:
  61. mock_exists.return_value = False
  62. self.assertEqual(self.ermine_ctl.status, (-2, 'InvalidState'))
  63. def testStatusReturnsRCAndStdout(self):
  64. with mock.patch.object(
  65. ErmineCtl, 'exists', new_callable=mock.PropertyMock) as mock_exists, \
  66. mock.patch.object(
  67. ErmineCtl, '_ExecuteCommandAsync') as mock_execute:
  68. mock_proc = mock.create_autospec(subprocess.Popen, instance=True)
  69. mock_proc.communicate.return_value = 'foo', 'stderr'
  70. mock_proc.returncode = 10
  71. mock_execute.return_value = mock_proc
  72. self.assertEqual(self.ermine_ctl.status, (10, 'foo'))
  73. def testStatusReturnsTimeoutState(self):
  74. with mock.patch.object(
  75. ErmineCtl, 'exists', new_callable=mock.PropertyMock) as mock_exists, \
  76. mock.patch.object(
  77. ErmineCtl, '_ExecuteCommandAsync') as mock_execute, \
  78. mock.patch.object(logging, 'warning') as _:
  79. mock_proc = mock.create_autospec(subprocess.Popen, instance=True)
  80. mock_proc.wait.side_effect = subprocess.TimeoutExpired(
  81. 'cmd', 'some timeout')
  82. mock_execute.return_value = mock_proc
  83. self.assertEqual(self.ermine_ctl.status, (-1, 'Timeout'))
  84. def testExecuteCommandRaisesErrorOnNonZeroReturn(self):
  85. with mock.patch.object(ErmineCtl, '_ExecuteCommandAsync') as mock_execute:
  86. mock_proc = mock.create_autospec(subprocess.Popen, instance=True)
  87. mock_proc.returncode = 42
  88. mock_proc.communicate.return_value = ('foo', 'bar')
  89. mock_execute.return_value = mock_proc
  90. self.assertRaises(RuntimeError, self.ermine_ctl._ExecuteCommand,
  91. ['some', 'command', '--args'])
  92. mock_execute.assert_called_once_with(['some', 'command', '--args'])
  93. mock_proc = mock.create_autospec(subprocess.Popen, instance=True)
  94. mock_proc.returncode = 0
  95. mock_proc.communicate.return_value = ('foo', 'bar')
  96. mock_execute.return_value = mock_proc
  97. self.assertIsNone(self.ermine_ctl._ExecuteCommand(['another', 'command']))
  98. def testExecuteCommandAsync(self):
  99. mock_proc = mock.create_autospec(subprocess.Popen, instance=True)
  100. self.mock_target.RunCommandPiped.return_value = mock_proc
  101. self.ermine_ctl._TOOL = 'foo'
  102. self.ermine_ctl._OOBE_SUBTOOL = 'oobe_buzz'
  103. self.assertEqual(
  104. self.ermine_ctl._ExecuteCommandAsync(['some', 'command', '--args']),
  105. mock_proc)
  106. self.mock_target.RunCommandPiped.assert_called_once_with(
  107. ['foo', 'oobe_buzz', 'some', 'command', '--args'],
  108. stdout=subprocess.PIPE,
  109. stderr=subprocess.STDOUT,
  110. text=True)
  111. self.assertEqual(self.mock_target.RunCommand.call_count, 0)
  112. def testWaitUntilReadyIsNotReadyIfDNE(self):
  113. with mock.patch.object(ErmineCtl, 'exists',
  114. new_callable=mock.PropertyMock) as mock_exists:
  115. mock_exists.return_value = False
  116. self.assertFalse(self.ermine_ctl.WaitUntilReady())
  117. def testWaitUntilReadyLoopsUntilReady(self):
  118. with mock.patch.object(ErmineCtl, 'exists',
  119. new_callable=mock.PropertyMock) as mock_exists, \
  120. mock.patch.object(time, 'sleep') as mock_sleep, \
  121. mock.patch.object(ErmineCtl, 'ready',
  122. new_callable=mock.PropertyMock) as mock_ready:
  123. mock_exists.return_value = True
  124. mock_ready.side_effect = [False, False, False, True]
  125. self.assertTrue(self.ermine_ctl.WaitUntilReady())
  126. self.assertEqual(mock_ready.call_count, 4)
  127. self.assertEqual(mock_sleep.call_count, 3)
  128. def testWaitUntilReadyLoopsUntilReady(self):
  129. with mock.patch.object(ErmineCtl, 'exists',
  130. new_callable=mock.PropertyMock) as mock_exists, \
  131. mock.patch.object(time, 'sleep') as mock_sleep, \
  132. mock.patch.object(ErmineCtl, 'ready',
  133. new_callable=mock.PropertyMock) as mock_ready:
  134. mock_exists.return_value = True
  135. mock_ready.side_effect = [False, False, False, True]
  136. self.ermine_ctl._WAIT_ATTEMPTS = 3
  137. self.assertRaises(TimeoutError, self.ermine_ctl.WaitUntilReady)
  138. self.assertEqual(mock_ready.call_count, 3)
  139. self.assertEqual(mock_sleep.call_count, 3)
  140. def testTakeToShellAssertsReady(self):
  141. with mock.patch.object(ErmineCtl, 'WaitUntilReady') as mock_wait_ready:
  142. mock_wait_ready.return_value = False
  143. self.assertRaises(AssertionError, self.ermine_ctl.TakeToShell)
  144. def testTakeToShellExitsOnCompleteState(self):
  145. with mock.patch.object(ErmineCtl, 'WaitUntilReady') as mock_wait_ready, \
  146. mock.patch.object(
  147. ErmineCtl, 'status',
  148. new_callable=mock.PropertyMock) as mock_status, \
  149. mock.patch.object(ErmineCtl, '_ExecuteCommand') as mock_execute:
  150. mock_wait_ready.return_value = True
  151. mock_status.return_value = (0, 'SomeCompleteState')
  152. self.ermine_ctl._COMPLETE_STATE = 'SomeCompleteState'
  153. self.ermine_ctl.TakeToShell()
  154. self.assertEqual(mock_execute.call_count, 0)
  155. def testTakeToShellInvalidStateRaisesException(self):
  156. with mock.patch.object(ErmineCtl, 'WaitUntilReady') as mock_wait_ready, \
  157. mock.patch.object(
  158. ErmineCtl, 'status',
  159. new_callable=mock.PropertyMock) as mock_status, \
  160. mock.patch.object(ErmineCtl, '_ExecuteCommand') as mock_execute:
  161. mock_wait_ready.return_value = True
  162. mock_status.return_value = (0, 'SomeUnknownState')
  163. self.assertNotIn('SomeUnknownState', self.ermine_ctl._STATE_TO_NEXT)
  164. self.assertRaises(NotImplementedError, self.ermine_ctl.TakeToShell)
  165. def testTakeToShellWithMaxTransitionsRaisesError(self):
  166. with mock.patch.object(ErmineCtl, 'WaitUntilReady') as mock_wait_ready, \
  167. mock.patch.object(
  168. ErmineCtl, 'status',
  169. new_callable=mock.PropertyMock) as mock_status, \
  170. mock.patch.object(ErmineCtl, '_ExecuteCommand') as mock_execute:
  171. mock_wait_ready.return_value = True
  172. self.ermine_ctl._MAX_STATE_TRANSITIONS = 3
  173. # Returns too many state transitions before CompleteState.
  174. mock_status.side_effect = [(0, 'Unknown'), (0, 'KnownWithPassword'),
  175. (0, 'Unknown'), (0, 'KnownWithPassword'),
  176. (0, 'CompleteState')]
  177. self.ermine_ctl._COMPLETE_STATE = 'CompleteState'
  178. self.ermine_ctl._STATE_TO_NEXT = {
  179. 'Unknown': 'fizz',
  180. 'KnownWithPassword': 'buzz',
  181. 'CompleteState': 'do nothing'
  182. }
  183. self.assertRaises(RuntimeError, self.ermine_ctl.TakeToShell)
  184. def testTakeToShellExecutesKnownCommands(self):
  185. with mock.patch.object(ErmineCtl, 'WaitUntilReady') as mock_wait_ready, \
  186. mock.patch.object(
  187. ErmineCtl, 'status',
  188. new_callable=mock.PropertyMock) as mock_status, \
  189. mock.patch.object(ErmineCtl, '_ExecuteCommand') as mock_execute:
  190. mock_wait_ready.return_value = True
  191. mock_status.side_effect = [(0, 'Unknown'), (0, 'KnownWithPassword'),
  192. (0, 'CompleteState')]
  193. self.ermine_ctl._COMPLETE_STATE = 'CompleteState'
  194. self.ermine_ctl._STATE_TO_NEXT = {
  195. 'Unknown': ['fizz'],
  196. 'KnownWithPassword': ['buzz', 'super plaintext password'],
  197. 'CompleteState': ['do nothing']
  198. }
  199. self.ermine_ctl.TakeToShell()
  200. self.assertEqual(mock_execute.call_count, 2)
  201. mock_execute.assert_has_calls([
  202. mock.call(['fizz']),
  203. mock.call(['buzz', 'super plaintext password'])
  204. ])
  205. if __name__ == '__main__':
  206. unittest.main()