src/third_party/chromite/cros/commands/cros_flash_unittest.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2013 The Chromium OS 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
   6 """This module tests the cros flash command."""
   7
   8 from __future__ import print_function
   9
  10 import mock
  11 import os
  12 import sys
  13
  14 sys.path.insert(0, os.path.abspath('%s/../../..' % os.path.dirname(__file__)))
  15 from chromite.cros.commands import cros_flash
  16 from chromite.cros.commands import init_unittest
  17 from chromite.lib import cros_build_lib
  18 from chromite.lib import cros_test_lib
  19 from chromite.lib import dev_server_wrapper
  20 from chromite.lib import partial_mock
  21 from chromite.lib import remote_access
  22
  23
  24 # pylint: disable=W0212
  25 class TestXbuddyHelpers(cros_test_lib.MockTempDirTestCase):
  26   """Test xbuddy helper functions."""
  27   def testGenerateXbuddyRequestForUpdate(self):
  28     """Test we generate correct xbuddy requests."""
  29     # Use the latest build when 'latest' is given.
  30     req = 'xbuddy/latest?for_update=true&return_dir=true'
  31     self.assertEqual(
  32         cros_flash.GenerateXbuddyRequest('latest', 'update'), req)
  33
  34     # Convert the path starting with 'xbuddy://' to 'xbuddy/'
  35     path = 'xbuddy://remote/stumpy/version'
  36     req = 'xbuddy/remote/stumpy/version?for_update=true&return_dir=true'
  37     self.assertEqual(cros_flash.GenerateXbuddyRequest(path, 'update'), req)
  38
  39   def testGenerateXbuddyRequestForImage(self):
  40     """Tests that we generate correct requests to get images."""
  41     image_path = 'foo/bar/taco'
  42     self.assertEqual(cros_flash.GenerateXbuddyRequest(image_path, 'image'),
  43                      'xbuddy/foo/bar/taco?return_dir=true')
  44
  45     image_path = 'xbuddy://foo/bar/taco'
  46     self.assertEqual(cros_flash.GenerateXbuddyRequest(image_path, 'image'),
  47                      'xbuddy/foo/bar/taco?return_dir=true')
  48
  49   def testGenerateXbuddyRequestForTranslate(self):
  50     """Tests that we generate correct requests for translation."""
  51     image_path = 'foo/bar/taco'
  52     self.assertEqual(cros_flash.GenerateXbuddyRequest(image_path, 'translate'),
  53                      'xbuddy_translate/foo/bar/taco')
  54
  55     image_path = 'xbuddy://foo/bar/taco'
  56     self.assertEqual(cros_flash.GenerateXbuddyRequest(image_path, 'translate'),
  57                      'xbuddy_translate/foo/bar/taco')
  58
  59   def testConvertTranslatedPath(self):
  60     """Tests that we convert a translated path to a usable xbuddy path."""
  61     path = 'remote/latest-canary'
  62     translated_path = 'taco-release/R36-5761.0.0/chromiumos_test_image.bin'
  63     self.assertEqual(cros_flash.ConvertTranslatedPath(path, translated_path),
  64                      'remote/taco-release/R36-5761.0.0/test')
  65
  66     path = 'latest'
  67     translated_path = 'taco/R36-5600.0.0/chromiumos_image.bin'
  68     self.assertEqual(cros_flash.ConvertTranslatedPath(path, translated_path),
  69                      'local/taco/R36-5600.0.0/dev')
  70
  71   @mock.patch('chromite.lib.cros_build_lib.IsInsideChroot', return_value=True)
  72   def testDevserverURLToLocalPath(self, _mock1):
  73     """Tests that we convert a devserver URL to a local path correctly."""
  74     url = 'http://localhost:8080/static/peppy-release/R33-5116.87.0'
  75     base_path = os.path.join(self.tempdir, 'peppy-release/R33-5116.87.0')
  76
  77     local_path = os.path.join(base_path, 'recovery_image.bin')
  78     self.assertEqual(
  79       cros_flash.DevserverURLToLocalPath(
  80           url, self.tempdir, 'recovery'), local_path)
  81
  82     # Default to test image.
  83     local_path = os.path.join(base_path, 'chromiumos_test_image.bin')
  84     self.assertEqual(
  85       cros_flash.DevserverURLToLocalPath(url, self.tempdir, 'taco'), local_path)
  86
  87
  88 class MockFlashCommand(init_unittest.MockCommand):
  89   """Mock out the flash command."""
  90   TARGET = 'chromite.cros.commands.cros_flash.FlashCommand'
  91   TARGET_CLASS = cros_flash.FlashCommand
  92   COMMAND = 'flash'
  93
  94   def __init__(self, *args, **kwargs):
  95     init_unittest.MockCommand.__init__(self, *args, **kwargs)
  96
  97   def Run(self, inst):
  98     init_unittest.MockCommand.Run(self, inst)
  99
 100
 101 class RemoteDeviceUpdaterMock(partial_mock.PartialCmdMock):
 102   """Mock out RemoteDeviceUpdater."""
 103   TARGET = 'chromite.cros.commands.cros_flash.RemoteDeviceUpdater'
 104   ATTRS = ('UpdateStateful', 'UpdateRootfs', 'GetUpdatePayloads',
 105            'SetupRootfsUpdate', 'Verify')
 106
 107   def __init__(self):
 108     partial_mock.PartialCmdMock.__init__(self)
 109
 110   def GetUpdatePayloads(self, _inst, *_args, **_kwargs):
 111     """Mock out GetUpdatePayloads."""
 112
 113   def UpdateStateful(self, _inst, *_args, **_kwargs):
 114     """Mock out UpdateStateful."""
 115
 116   def UpdateRootfs(self, _inst, *_args, **_kwargs):
 117     """Mock out UpdateRootfs."""
 118
 119   def SetupRootfsUpdate(self, _inst, *_args, **_kwargs):
 120     """Mock out SetupRootfsUpdate."""
 121
 122   def Verify(self, _inst, *_args, **_kwargs):
 123     """Mock out SetupRootfsUpdate."""
 124
 125
 126 class UpdateRunThroughTest(cros_test_lib.MockTempDirTestCase,
 127                            cros_test_lib.LoggingTestCase):
 128   """Test the flow of FlashCommand.run with the update methods mocked out."""
 129
 130   IMAGE = '/path/to/image'
 131   DEVICE = '1.1.1.1'
 132
 133   def SetupCommandMock(self, cmd_args):
 134     """Setup comand mock."""
 135     self.cmd_mock = MockFlashCommand(
 136         cmd_args, base_args=['--cache-dir', self.tempdir])
 137     self.StartPatcher(self.cmd_mock)
 138
 139   def setUp(self):
 140     """Patches objects."""
 141     self.cmd_mock = None
 142     self.updater_mock = self.StartPatcher(RemoteDeviceUpdaterMock())
 143     self.PatchObject(cros_flash, 'GenerateXbuddyRequest',
 144                      return_value='xbuddy/local/latest')
 145     self.PatchObject(dev_server_wrapper, 'DevServerWrapper')
 146     self.PatchObject(cros_flash, 'TranslateImagePath',
 147                      return_value='taco-paladin/R36/chromiumos_test_image.bin')
 148     self.PatchObject(remote_access, 'CHECK_INTERVAL', new=0)
 149     self.PatchObject(remote_access.ChromiumOSDevice, '_LearnBoard',
 150                      return_value='peppy')
 151
 152   def testUpdateAll(self):
 153     """Tests that update methods are called correctly."""
 154     self.SetupCommandMock([self.DEVICE, self.IMAGE])
 155     with mock.patch('os.path.exists', return_value=True):
 156       self.cmd_mock.inst.Run()
 157       self.assertTrue(self.updater_mock.patched['UpdateStateful'].called)
 158       self.assertTrue(self.updater_mock.patched['UpdateRootfs'].called)
 159
 160   def testUpdateStateful(self):
 161     """Tests that update methods are called correctly."""
 162     self.SetupCommandMock(['--no-rootfs-update', self.DEVICE, self.IMAGE])
 163     with mock.patch('os.path.exists', return_value=True):
 164       self.cmd_mock.inst.Run()
 165       self.assertTrue(self.updater_mock.patched['UpdateStateful'].called)
 166       self.assertFalse(self.updater_mock.patched['UpdateRootfs'].called)
 167
 168   def testUpdateRootfs(self):
 169     """Tests that update methods are called correctly."""
 170     self.SetupCommandMock(['--no-stateful-update', self.DEVICE, self.IMAGE])
 171     with mock.patch('os.path.exists', return_value=True):
 172       self.cmd_mock.inst.Run()
 173       self.assertFalse(self.updater_mock.patched['UpdateStateful'].called)
 174       self.assertTrue(self.updater_mock.patched['UpdateRootfs'].called)
 175
 176   def testMissingPayloads(self):
 177     """Tests we exit when payloads are missing."""
 178     self.SetupCommandMock([self.DEVICE, self.IMAGE])
 179     with mock.patch('os.path.exists', return_value=False):
 180       self.assertRaises(cros_build_lib.DieSystemExit, self.cmd_mock.inst.Run)
 181
 182
 183 class USBImagerMock(partial_mock.PartialCmdMock):
 184   """Mock out USBImager."""
 185   TARGET = 'chromite.cros.commands.cros_flash.USBImager'
 186   ATTRS = ('GetImagePathFromDevserver', 'CopyImageToDevice',
 187            'InstallImageToDevice', 'ChooseRemovableDevice',
 188            'ListAllRemovableDevices', 'GetRemovableDeviceDescription',
 189            'IsFilePathGPTDiskImage')
 190   VALID_IMAGE = True
 191
 192   def __init__(self):
 193     partial_mock.PartialCmdMock.__init__(self)
 194
 195   def GetImagePathFromDevserver(self, _inst, *_args, **_kwargs):
 196     """Mock out GetImagePathFromDevserver."""
 197
 198   def CopyImageToDevice(self, _inst, *_args, **_kwargs):
 199     """Mock out CopyImageToDevice."""
 200
 201   def InstallImageToDevice(self, _inst, *_args, **_kwargs):
 202     """Mock out InstallImageToDevice."""
 203
 204   def ChooseRemovableDevice(self, _inst, *_args, **_kwargs):
 205     """Mock out ChooseRemovableDevice."""
 206
 207   def ListAllRemovableDevices(self, _inst, *_args, **_kwargs):
 208     """Mock out ListAllRemovableDevices."""
 209     return ['foo', 'taco', 'milk']
 210
 211   def GetRemovableDeviceDescription(self, _inst, *_args, **_kwargs):
 212     """Mock out GetRemovableDeviceDescription."""
 213
 214   def IsFilePathGPTDiskImage(self, _inst, *_args, **_kwargs):
 215     """Mock out IsFilePathGPTDiskImage."""
 216     return self.VALID_IMAGE
 217
 218
 219 class ImagingRunThroughTest(cros_test_lib.MockTempDirTestCase,
 220                             cros_test_lib.LoggingTestCase):
 221   """Test the flow of FlashCommand.run with the imaging methods mocked out."""
 222   IMAGE = '/path/to/image'
 223
 224   def SetupCommandMock(self, cmd_args):
 225     """Setup comand mock."""
 226     self.cmd_mock = MockFlashCommand(
 227         cmd_args, base_args=['--cache-dir', self.tempdir])
 228     self.StartPatcher(self.cmd_mock)
 229
 230   def setUp(self):
 231     """Patches objects."""
 232     self.cmd_mock = None
 233     self.usb_mock = USBImagerMock()
 234     self.imager_mock = self.StartPatcher(self.usb_mock)
 235     self.PatchObject(cros_flash, 'GenerateXbuddyRequest',
 236                      return_value='xbuddy/local/latest')
 237     self.PatchObject(dev_server_wrapper, 'DevServerWrapper')
 238     self.PatchObject(cros_flash, 'TranslateImagePath',
 239                      return_value='taco-paladin/R36/chromiumos_test_image.bin')
 240     self.PatchObject(os.path, 'exists', return_value=True)
 241
 242   def testLocalImagePathCopy(self):
 243     """Tests that imaging methods are called correctly."""
 244     self.SetupCommandMock(['usb:///dev/foo', self.IMAGE])
 245     with mock.patch('os.path.isfile', return_value=True):
 246       self.cmd_mock.inst.Run()
 247       self.assertTrue(self.imager_mock.patched['CopyImageToDevice'].called)
 248
 249   def testLocalImagePathInstall(self):
 250     """Tests that imaging methods are called correctly."""
 251     self.SetupCommandMock(['--board=taco', '--install', 'usb:///dev/foo',
 252                            self.IMAGE])
 253     with mock.patch('os.path.isfile', return_value=True):
 254       self.cmd_mock.inst.Run()
 255       self.assertTrue(self.imager_mock.patched['InstallImageToDevice'].called)
 256
 257   def testLocalBadImagePath(self):
 258     """Tests that using an image not having the magic bytes has prompt."""
 259     self.usb_mock.VALID_IMAGE = False
 260     self.SetupCommandMock(['usb:///dev/foo', self.IMAGE])
 261     with mock.patch('os.path.isfile', return_value=True):
 262       with mock.patch.object(cros_build_lib, 'BooleanPrompt') as mock_prompt:
 263         mock_prompt.return_value = False
 264         self.cmd_mock.inst.Run()
 265         self.assertTrue(mock_prompt.called)
 266
 267   def testNonLocalImgePath(self):
 268     """Tests that we try to get the image path from devserver."""
 269     self.SetupCommandMock(['usb:///dev/foo', self.IMAGE])
 270     with mock.patch('os.path.isfile', return_value=False):
 271       with mock.patch('os.path.isdir', return_value=False):
 272         self.cmd_mock.inst.Run()
 273         self.assertTrue(
 274             self.imager_mock.patched['GetImagePathFromDevserver'].called)
 275         self.assertTrue(self.imager_mock.patched['CopyImageToDevice'].called)
 276
 277   def testConfirmNonRemovableDevice(self):
 278     """Tests that we ask user to confirm if the device is not removable."""
 279     with mock.patch.object(cros_build_lib, 'BooleanPrompt') as mock_prompt:
 280       self.SetupCommandMock(['usb:///dev/dummy', self.IMAGE])
 281       self.cmd_mock.inst.Run()
 282       self.assertTrue(mock_prompt.called)
 283
 284   def testSkipPromptNonRemovableDevice(self):
 285     """Tests that we skip the prompt for non-removable with --yes."""
 286     with mock.patch.object(cros_build_lib, 'BooleanPrompt') as mock_prompt:
 287       self.SetupCommandMock(['--yes', 'usb:///dev/dummy', self.IMAGE])
 288       self.cmd_mock.inst.Run()
 289       self.assertFalse(mock_prompt.called)
 290
 291   def testChooseRemovableDevice(self):
 292     """Tests that we ask user to choose a device if none is given."""
 293     self.SetupCommandMock(['usb://', self.IMAGE])
 294     self.cmd_mock.inst.Run()
 295     self.assertTrue(self.imager_mock.patched['ChooseRemovableDevice'].called)
 296
 297
 298 if __name__ == '__main__':
 299   cros_test_lib.main()