Home | History | Annotate | Download | only in host_controller
      1 #!/usr/bin/env python
      2 #
      3 # Copyright (C) 2017 The Android Open Source Project
      4 #
      5 # Licensed under the Apache License, Version 2.0 (the "License");
      6 # you may not use this file except in compliance with the License.
      7 # You may obtain a copy of the License at
      8 #
      9 #      http://www.apache.org/licenses/LICENSE-2.0
     10 #
     11 # Unless required by applicable law or agreed to in writing, software
     12 # distributed under the License is distributed on an "AS IS" BASIS,
     13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     14 # See the License for the specific language governing permissions and
     15 # limitations under the License.
     16 #
     17 
     18 import os
     19 import unittest
     20 
     21 try:
     22     from unittest import mock
     23 except ImportError:
     24     import mock
     25 
     26 try:
     27     import StringIO as string_io_module
     28 except ImportError:
     29     import io as string_io_module
     30 
     31 from host_controller.build import build_flasher
     32 from host_controller.tfc import command_task
     33 from host_controller.tfc import device_info
     34 from host_controller import console
     35 
     36 
     37 class ConsoleTest(unittest.TestCase):
     38     """A test for console.Console.
     39 
     40     Attribute:
     41         _out_file: The console output buffer.
     42         _host_controller: A mock tfc_host_controller.HostController.
     43         _build_provider_pab: A mock build_provider_pab.BuildProviderPAB.
     44         _tfc_client: A mock tfc_client.TfcClient.
     45         _vti_client A mock vti_endpoint_client.VtiEndpointClient.
     46         _console: The console being tested.
     47     """
     48     _DEVICES = [
     49         device_info.DeviceInfo(
     50             device_serial="ABC001",
     51             run_target="sailfish",
     52             state="Available",
     53             build_id="111111",
     54             sdk_version="27")
     55     ]
     56     _TASKS = [
     57         command_task.CommandTask(
     58             request_id="1",
     59             task_id="1-0",
     60             command_id="2",
     61             command_line="vts -m SampleShellTest",
     62             device_serials=["ABC001"])
     63     ]
     64 
     65     def setUp(self):
     66         """Creates the console."""
     67         self._out_file = string_io_module.StringIO()
     68         self._host_controller = mock.Mock()
     69         self._build_provider_pab = mock.Mock()
     70         self._tfc_client = mock.Mock()
     71         self._vti_client = mock.Mock()
     72         self._console = console.Console(
     73             self._vti_client, self._tfc_client, self._build_provider_pab,
     74             [self._host_controller],
     75             None, self._out_file)
     76 
     77     def tearDown(self):
     78         """Closes the output file."""
     79         self._out_file.close()
     80 
     81     def _IssueCommand(self, command_line):
     82         """Issues a command in the console.
     83 
     84         Args:
     85             command_line: A string, the input to the console.
     86 
     87         Returns:
     88             A string, the output of the console.
     89         """
     90         out_position = self._out_file.tell()
     91         self._console.onecmd(command_line)
     92         self._out_file.seek(out_position)
     93         return self._out_file.read()
     94 
     95     def testLease(self):
     96         """Tests the lease command."""
     97         self._host_controller.LeaseCommandTasks.return_value = self._TASKS
     98         output = self._IssueCommand("lease")
     99         expected = (
    100             "request_id  command_id  task_id  device_serials  command_line          \n"
    101             "1           2           1-0      ABC001          vts -m SampleShellTest\n"
    102         )
    103         self.assertEqual(expected, output)
    104         output = self._IssueCommand("lease --host 0")
    105         self.assertEqual(expected, output)
    106 
    107     def testRequest(self):
    108         """Tests the request command."""
    109         user = "user0"
    110         cluster = "cluster0"
    111         run_target = "sailfish"
    112         command_line = "vts -m SampleShellTest"
    113         self._IssueCommand("request --user %s --cluster %s --run-target %s "
    114                            "-- %s" % (user, cluster, run_target, command_line))
    115         req = self._tfc_client.NewRequest.call_args[0][0]
    116         self.assertEqual(user, req.user)
    117         self.assertEqual(cluster, req.cluster)
    118         self.assertEqual(run_target, req.run_target)
    119         self.assertEqual(command_line, req.command_line)
    120 
    121     def testListHosts(self):
    122         """Tests the list command."""
    123         self._host_controller.hostname = "host0"
    124         output = self._IssueCommand("list hosts")
    125         self.assertEqual("index  name\n" "[  0]  host0\n", output)
    126 
    127     def testListDevices(self):
    128         """Tests the list command."""
    129         self._host_controller.ListDevices.return_value = self._DEVICES
    130         self._host_controller.hostname = "host0"
    131         output = self._IssueCommand("list devices")
    132         expected = (
    133             "[  0]  host0\n"
    134             "device_serial  state      run_target  build_id  sdk_version  stub\n"
    135             "ABC001         Available  sailfish    111111    27               \n"
    136         )
    137         self.assertEqual(expected, output)
    138         output = self._IssueCommand("list devices --host 0")
    139         self.assertEqual(expected, output)
    140 
    141     def testWrongHostIndex(self):
    142         """Tests host index out of range."""
    143         output = self._IssueCommand("list devices --host 1")
    144         expected = "IndexError: "
    145         self.assertTrue(output.startswith(expected))
    146         output = self._IssueCommand("lease --host 1")
    147         self.assertTrue(output.startswith(expected))
    148 
    149     @mock.patch('host_controller.console.build_flasher.BuildFlasher')
    150     def testFetchPOSTAndFlash(self, mock_class):
    151         """Tests fetching from pab and flashing."""
    152         self._build_provider_pab.GetArtifact.return_value = ({
    153             "system.img":
    154             "/mock/system.img",
    155             "odm.img":
    156             "/mock/odm.img"
    157         }, {}, {"build_id": "build_id"}, {})
    158         self._IssueCommand(
    159             "fetch --branch=aosp-master-ndk --target=darwin_mac "
    160             "--account_id=100621237 "
    161             "--artifact_name=foo-{id}.tar.bz2 --method=POST")
    162         self._build_provider_pab.GetArtifact.assert_called_with(
    163             account_id='100621237',
    164             branch='aosp-master-ndk',
    165             target='darwin_mac',
    166             artifact_name='foo-{id}.tar.bz2',
    167             build_id='latest',
    168             method='POST')
    169 
    170         flasher = mock.Mock()
    171         mock_class.return_value = flasher
    172         self._IssueCommand("flash --current system=system.img odm=odm.img")
    173         flasher.Flash.assert_called_with({
    174             "system": "/mock/system.img",
    175             "odm": "/mock/odm.img"
    176         })
    177 
    178     def testFetchAndEnvironment(self):
    179         """Tests fetching from pab and check stored os environment"""
    180         build_id_return = "4328532"
    181         target_return = "darwin_mac"
    182         expected_fetch_info = {"build_id": build_id_return}
    183 
    184         self._build_provider_pab.GetArtifact.return_value = ({
    185             "system.img":
    186             "/mock/system.img",
    187             "odm.img":
    188             "/mock/odm.img"
    189         }, {}, expected_fetch_info, {})
    190         self._IssueCommand("fetch --branch=aosp-master-ndk --target=%s "
    191                            "--account_id=100621237 "
    192                            "--artifact_name=foo-{id}.tar.bz2 --method=POST"
    193                            % target_return)
    194         self._build_provider_pab.GetArtifact.assert_called_with(
    195             account_id='100621237',
    196             branch='aosp-master-ndk',
    197             target='darwin_mac',
    198             artifact_name='foo-{id}.tar.bz2',
    199             build_id='latest',
    200             method='POST')
    201 
    202         expected = expected_fetch_info["build_id"]
    203         self.assertEqual(build_id_return, expected)
    204 
    205     @mock.patch('host_controller.console.build_flasher.BuildFlasher')
    206     def testFlashGSI(self, mock_class):
    207         flasher = mock.Mock()
    208         mock_class.return_value = flasher
    209         self._IssueCommand("flash --gsi=system.img")
    210         flasher.FlashGSI.assert_called_with('system.img', None)
    211 
    212     @mock.patch('host_controller.console.build_flasher.BuildFlasher')
    213     def testFlashGSIWithVbmeta(self, mock_class):
    214         flasher = mock.Mock()
    215         mock_class.return_value = flasher
    216         self._IssueCommand("flash --gsi=system.img --vbmeta=vbmeta.img")
    217         flasher.FlashGSI.assert_called_with('system.img', 'vbmeta.img')
    218 
    219     @mock.patch('host_controller.console.build_flasher.BuildFlasher')
    220     def testFlashall(self, mock_class):
    221         flasher = mock.Mock()
    222         mock_class.return_value = flasher
    223         self._IssueCommand("flash --build_dir=path/to/dir/")
    224         flasher.Flashall.assert_called_with('path/to/dir/')
    225 
    226     @mock.patch('host_controller.console.importlib')
    227     @mock.patch('host_controller.console.issubclass')
    228     def testImportFlasher(self, mock_issubclass, mock_importlib):
    229         mock_issubclass.return_value = True
    230         flasher_module = mock.Mock()
    231         flasher = mock.Mock()
    232         mock_importlib.import_module.return_value = flasher_module
    233         flasher_module.Flasher.return_value = flasher
    234         self._IssueCommand("flash --serial ABC001 "
    235                            "--flasher_type test.flasher.Flasher "
    236                            "--flasher_path /test/flasher "
    237                            "-- --unit test")
    238         mock_issubclass.assert_called_once_with(
    239             flasher_module.Flasher, build_flasher.BuildFlasher)
    240         mock_importlib.import_module.assert_called_with("test.flasher")
    241         flasher_module.Flasher.assert_called_with("ABC001", "/test/flasher")
    242         flasher.Flash.assert_called_with({}, {}, "--unit", "test")
    243         flasher.WaitForDevice.assert_called_with()
    244 
    245 
    246 if __name__ == "__main__":
    247     unittest.main()
    248