1# Copyright (C) 2016 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from __future__ import print_function
16
17import argparse
18import os
19import unittest
20import fastboot
21import subprocess
22import sys
23
24# Default values for arguments
25device_type = "phone"
26
27class ShellTest(unittest.TestCase):
28    @classmethod
29    def setUpClass(cls):
30        cls.fastboot = fastboot.FastbootDevice()
31
32    def exists_validvals(self, varname, varlist, validlist):
33        self.assertIn(varname, varlist)
34        self.assertIn(varlist[varname], validlist)
35        return varlist[varname]
36
37    def exists_yes_no(self, varname, varlist):
38        return self.exists_validvals(varname, varlist, ["yes", "no"])
39
40    def exists_nonempty(self, varname, varlist):
41        self.assertIn(varname, varlist)
42        self.assertGreater(len(varlist[varname]), 0)
43        return varlist[varname]
44
45    def exists_integer(self, varname, varlist, base=10):
46        val = 0
47        self.assertIn(varname, varlist)
48        try:
49            val = int(varlist[varname], base)
50        except ValueError:
51            self.fail("%s (%s) is not an integer" % (varname, varlist[varname]))
52        return val
53
54    def get_exists(self, varname):
55        val = self.fastboot.getvar(varname)
56        self.assertIsNotNone(val)
57        return val
58
59    def get_exists_validvals(self, varname, validlist):
60        val = self.get_exists(varname)
61        self.assertIn(val, validlist)
62        return val
63
64    def get_exists_yes_no(self, varname):
65        return self.get_exists_validvals(varname, ["yes", "no"])
66
67    def get_exists_nonempty(self, varname):
68        val = self.get_exists(varname)
69        self.assertGreater(len(val), 0)
70        return val
71
72    def get_exists_integer(self, varname, base=10):
73        val = self.get_exists(varname)
74        try:
75            num = int(val, base)
76        except ValueError:
77            self.fail("%s (%s) is not an integer" % (varname, val))
78        return num
79
80    def get_slotcount(self):
81        slotcount = 0
82        try:
83            val = self.fastboot.getvar("slot-count")
84            if val != None:
85                slotcount = int(val)
86        except ValueError:
87            self.fail("slot-count (%s) is not an integer" % val)
88        except subprocess.CalledProcessError:
89            print("Does not appear to be an A/B device.")
90        if not slotcount:
91            print("Does not appear to be an A/B device.")
92        return slotcount
93
94    def test_getvarall(self):
95        """Tests that required variables are reported by getvar all"""
96
97        var_all = self.fastboot.getvar_all()
98        self.exists_nonempty("version-baseband", var_all)
99        self.exists_nonempty("version-bootloader", var_all)
100        self.exists_nonempty("product", var_all)
101        self.exists_yes_no("secure", var_all)
102        self.exists_yes_no("unlocked", var_all)
103        self.exists_validvals("off-mode-charge", var_all, ["0", "1"])
104        self.assertIn("variant", var_all)
105        voltage = self.exists_nonempty("battery-voltage", var_all)
106        if voltage[-2:].lower() == "mv":
107            voltage = voltage[:-2]
108        try:
109            voltnum = float(voltage)
110        except ValueError:
111            self.fail("battery-voltage (%s) is not a number" % (varname, voltage))
112        self.exists_yes_no("battery-soc-ok", var_all)
113        maxdl = self.exists_integer("max-download-size", var_all, 16)
114        self.assertGreater(maxdl, 0)
115
116        if "slot-count" in var_all:
117            try:
118                slotcount = int(var_all["slot-count"])
119            except ValueError:
120                self.fail("slot-count (%s) is not an integer" % var_all["slot-count"])
121            if slotcount > 1:
122                # test for A/B variables
123                slots = [chr(slotnum+ord('a')) for slotnum in range(slotcount)]
124                self.exists_validvals("current-slot", var_all, slots)
125
126                # test for slot metadata
127                for slot in slots:
128                    self.exists_yes_no("slot-unbootable:"+slot, var_all)
129                    self.exists_yes_no("slot-unbootable:"+slot, var_all)
130                    self.exists_integer("slot-retry-count:"+slot, var_all)
131            else:
132                print("This does not appear to be an A/B device.")
133
134    def test_getvar_nonexistent(self):
135        """Tests behaviour of nonexistent variables."""
136
137        self.assertIsNone(self.fastboot.getvar("fhqwhgads"))
138
139    def test_getvar(self):
140        """Tests all variables separately"""
141
142        self.get_exists_nonempty("version-baseband")
143        self.get_exists_nonempty("version-bootloader")
144        self.get_exists_nonempty("product")
145        self.get_exists_yes_no("secure")
146        self.get_exists_yes_no("unlocked")
147        self.get_exists_validvals("off-mode-charge", ["0", "1"])
148        self.get_exists("variant")
149        voltage = self.get_exists_nonempty("battery-voltage")
150        if voltage[-2:].lower() == "mv":
151            voltage = voltage[:-2]
152        try:
153            voltnum = float(voltage)
154        except ValueError:
155            self.fail("battery-voltage (%s) is not a number" % voltage)
156        self.get_exists_yes_no("battery-soc-ok")
157        maxdl = self.get_exists_integer("max-download-size", 16)
158        self.assertGreater(maxdl, 0)
159
160        slotcount = self.get_slotcount()
161        if slotcount  > 1:
162            # test for A/B variables
163            slots = [chr(slotnum+ord('a')) for slotnum in range(slotcount)]
164            self.get_exists_validvals("current-slot", slots)
165
166            # test for slot metadata
167            for slot in slots:
168                self.get_exists_yes_no("slot-unbootable:"+slot)
169                self.get_exists_yes_no("slot-successful:"+slot)
170                self.get_exists_integer("slot-retry-count:"+slot)
171
172    def test_setactive(self):
173        """Tests that A/B devices can switch to each slot, and the change persists over a reboot."""
174        # Test invalid if not an A/B device
175        slotcount = self.get_slotcount()
176        if not slotcount:
177            return
178
179        maxtries = 0
180        slots = [chr(slotnum+ord('a')) for slotnum in range(slotcount)]
181        for slot in slots:
182            self.fastboot.set_active(slot)
183            self.assertEqual(slot, self.fastboot.getvar("current-slot"))
184            self.assertEqual("no", self.fastboot.getvar("slot-unbootable:"+slot))
185            self.assertEqual("no", self.fastboot.getvar("slot-successful:"+slot))
186            retry = self.get_exists_integer("slot-retry-count:"+slot)
187            if maxtries == 0:
188                maxtries = retry
189            else:
190                self.assertEqual(maxtries, retry)
191            self.fastboot.reboot(True)
192            self.assertEqual(slot, self.fastboot.getvar("current-slot"))
193            self.assertEqual("no", self.fastboot.getvar("slot-unbootable:"+slot))
194            self.assertEqual("no", self.fastboot.getvar("slot-successful:"+slot))
195            retry = self.get_exists_integer("slot-retry-count:"+slot)
196            if maxtries == 0:
197                maxtries = retry
198            else:
199                self.assertEqual(maxtries, retry)
200
201    def test_hasslot(self):
202        """Tests that A/B devices report partitions that have slots."""
203        # Test invalid if not an A/B device
204        if not self.get_slotcount():
205            return
206
207        self.assertEqual("yes", self.fastboot.getvar("has-slot:system"))
208        self.assertEqual("yes", self.fastboot.getvar("has-slot:boot"))
209
210        # Additional partition on AndroidThings (IoT) devices
211        if device_type == "iot":
212            self.assertEqual("yes", self.fastboot.getvar("has-slot:oem"))
213
214if __name__ == '__main__':
215    parser = argparse.ArgumentParser()
216    parser.add_argument("--device-type", default="phone",
217                        help="Type of device ('phone' or 'iot').")
218    parser.add_argument("extra_args", nargs="*")
219    args = parser.parse_args()
220
221    if args.device_type.lower() not in ("phone", "iot"):
222        raise ValueError("Unsupported device type '%s'." % args.device_type)
223    device_type = args.device_type.lower()
224
225    sys.argv[1:] = args.extra_args
226    unittest.main(verbosity=3)
227