1# Copyright 2018 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Test the install module."""
6
7import mock
8import unittest
9
10import common
11from autotest_lib.site_utils.deployment import install
12from autotest_lib.site_utils.stable_images import assign_stable_images
13
14
15class AFEMock(object):
16    """Mock frontend.AFE."""
17    CROS_IMAGE_TYPE = 'cros'
18    FIRMWARE_IMAGE_TYPE = 'firmware'
19
20    def __init__(self, cros_version=None, fw_version=None):
21        self.cros_version_map = mock.Mock()
22        self.cros_version_map.get_version.return_value = cros_version
23        self.fw_version_map = mock.Mock()
24        self.fw_version_map.get_version.return_value = fw_version
25
26    def get_stable_version_map(self, image_type):
27        if image_type == self.CROS_IMAGE_TYPE:
28            return self.cros_version_map
29        elif image_type == self.FIRMWARE_IMAGE_TYPE:
30             return self.fw_version_map
31
32
33class UpdateBuildTests(unittest.TestCase):
34    """Tests for _update_build."""
35
36    OMAHA_VERSION = 'R64-10176.65.0'
37    CROS_VERSION = 'R64-10175.65.0'
38
39    WOLF_BOARD = 'wolf'
40    WOLF_FW_VERSION = 'Google_Wolf.4389.24.62'
41    WOLF_NEW_FW_VERSION = 'Google_Wolf.4390.24.62'
42    WOLF_FW_VERSION_MAP = {WOLF_BOARD: WOLF_NEW_FW_VERSION}
43
44    CORAL_BOARD = 'coral'
45    CORAL_FW_VERSION = 'Google_Coral.10068.37.0'
46    CORAL_FW_VERSION_MAP = {
47            'blue': 'Google_Coral.10068.39.0',
48            'robo360': 'Google_Coral.10068.34.0',
49            'porbeagle': None,
50    }
51
52    def setUp(self):
53        self.report_log_mock = mock.Mock()
54        self.patchers = []
55
56    def _set_patchers(self,
57                      omaha_version=OMAHA_VERSION,
58                      firmware_versions=WOLF_FW_VERSION_MAP):
59        patcher1 = mock.patch.object(
60                install, '_get_omaha_build',
61                return_value=omaha_version)
62        patcher2 = mock.patch.object(
63                assign_stable_images, 'get_firmware_versions',
64                return_value=firmware_versions)
65
66        self.patchers.extend([patcher1, patcher2])
67
68        for p in self.patchers:
69            p.start()
70
71    def tearDown(self):
72        for p in self.patchers:
73            p.stop()
74
75    def test_update_build_cros_and_fw_version_on_non_unibuild(self):
76        """Update non-unibuild with old cros_version and fw_version in AFE."""
77        afe_mock = AFEMock(
78                cros_version=self.CROS_VERSION, fw_version=self.WOLF_FW_VERSION)
79
80        self._set_patchers()
81        arguments = mock.Mock(board=self.WOLF_BOARD, nostable=False, build=None)
82        cros_version = install._update_build(
83                afe_mock, self.report_log_mock, arguments)
84
85        afe_mock.cros_version_map.set_version.assert_called_once_with(
86                self.WOLF_BOARD, self.OMAHA_VERSION)
87        afe_mock.fw_version_map.set_version.assert_called_once_with(
88                self.WOLF_BOARD, self.WOLF_NEW_FW_VERSION)
89        self.assertEqual(cros_version, self.OMAHA_VERSION)
90
91    def test_update_build_without_omaha_version_on_non_unibuild(self):
92        """Do not update non-unibuild as no OMAHA_VERSION found."""
93        afe_mock = AFEMock(
94                cros_version=self.CROS_VERSION, fw_version=self.WOLF_FW_VERSION)
95
96        self._set_patchers(omaha_version=None)
97        arguments = mock.Mock(board=self.WOLF_BOARD, nostable=False, build=None)
98        cros_version = install._update_build(
99                afe_mock, self.report_log_mock, arguments)
100
101        afe_mock.cros_version_map.set_version.assert_not_called()
102        afe_mock.cros_version_map.delete_version.assert_not_called()
103        afe_mock.fw_version_map.set_version.assert_not_called()
104        afe_mock.fw_version_map.delete_version.assert_not_called()
105        self.assertEqual(cros_version, self.CROS_VERSION)
106
107    def test_update_build_cros_on_non_unibuild(self):
108        """Update non-unibuild with old cros_version in AFE."""
109        afe_mock = AFEMock(
110                cros_version=self.CROS_VERSION, fw_version=self.WOLF_FW_VERSION)
111        self._set_patchers(
112                firmware_versions={self.WOLF_BOARD: self.WOLF_FW_VERSION})
113        arguments = mock.Mock(board=self.WOLF_BOARD, nostable=False, build=None)
114        cros_version = install._update_build(
115                afe_mock, self.report_log_mock, arguments)
116
117        afe_mock.cros_version_map.set_version.assert_called_once_with(
118                self.WOLF_BOARD, self.OMAHA_VERSION)
119        afe_mock.fw_version_map.set_version.assert_not_called()
120        afe_mock.fw_version_map.delete_version.assert_not_called()
121        self.assertEqual(cros_version, self.OMAHA_VERSION)
122
123    def test_update_build_none_cros_and_fw_version_on_non_unibuild(self):
124        """Update Non-unibuild with None cros_version & fw_version in AFE."""
125        afe_mock = AFEMock(cros_version=None, fw_version=None)
126        self._set_patchers()
127        arguments = mock.Mock(board=self.WOLF_BOARD, nostable=False, build=None)
128        cros_version = install._update_build(
129                afe_mock, self.report_log_mock, arguments)
130
131        afe_mock.cros_version_map.set_version.assert_called_once_with(
132                self.WOLF_BOARD, self.OMAHA_VERSION)
133        afe_mock.fw_version_map.set_version.assert_called_once_with(
134                self.WOLF_BOARD, self.WOLF_NEW_FW_VERSION)
135        self.assertEqual(cros_version, self.OMAHA_VERSION)
136
137    def test_update_build_cros_and_fw_version_on_unibuild(self):
138        """Update unibuild with old cros_version and fw_versions."""
139        afe_mock = AFEMock(
140                cros_version=self.CROS_VERSION,
141                fw_version=self.CORAL_FW_VERSION)
142        self._set_patchers(
143                firmware_versions=self.CORAL_FW_VERSION_MAP)
144        arguments = mock.Mock(board=self.CORAL_BOARD, nostable=False,
145                              build=None)
146        cros_version = install._update_build(
147                afe_mock, self.report_log_mock, arguments)
148
149        afe_mock.cros_version_map.set_version.assert_called_once_with(
150                self.CORAL_BOARD, self.OMAHA_VERSION)
151        afe_mock.fw_version_map.set_version.assert_any_call(
152                'blue', 'Google_Coral.10068.39.0')
153        afe_mock.fw_version_map.set_version.assert_any_call(
154                'robo360', 'Google_Coral.10068.34.0')
155        afe_mock.fw_version_map.delete_version.assert_any_call(
156                'porbeagle')
157        self.assertEqual(cros_version, self.OMAHA_VERSION)
158
159
160if __name__ == '__main__':
161    unittest.main()
162