# Copyright (C) 2016 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import adb
import os
import unittest
import fastboot
import subprocess

class ShellTest(unittest.TestCase):
    def __init__(self, *args, **kwargs):
        super(ShellTest, self).__init__(*args, **kwargs)
        self.fastboot = fastboot.FastbootDevice()

    def exists_validvals(self, varname, varlist, validlist):
        self.assertIn(varname, varlist)
        self.assertIn(varlist[varname], validlist)
        return varlist[varname]

    def exists_yes_no(self, varname, varlist):
        return self.exists_validvals(varname, varlist, ["yes", "no"])

    def exists_nonempty(self, varname, varlist):
        self.assertIn(varname, varlist)
        self.assertGreater(len(varlist[varname]), 0)
        return varlist[varname]

    def exists_integer(self, varname, varlist, base=10):
        val = 0
        self.assertIn(varname, varlist)
        try:
            val = int(varlist[varname], base)
        except ValueError:
            self.fail("%s (%s) is not an integer" % (varname, varlist[varname]))
        return val

    def get_exists(self, varname):
        val = self.fastboot.getvar(varname)
        self.assertIsNotNone(val)
        return val

    def get_exists_validvals(self, varname, validlist):
        val = self.get_exists(varname)
        self.assertIn(val, validlist)
        return val

    def get_exists_yes_no(self, varname):
        return self.get_exists_validvals(varname, ["yes", "no"])

    def get_exists_nonempty(self, varname):
        val = self.get_exists(varname)
        self.assertGreater(len(val), 0)
        return val

    def get_exists_integer(self, varname, base=10):
        val = self.get_exists(varname)
        try:
            num = int(val, base)
        except ValueError:
            self.fail("%s (%s) is not an integer" % (varname, val))
        return num

    def test_getvarall(self):
        """Tests that required variables are reported by getvar all"""

        var_all = self.fastboot.getvar_all()
        self.exists_nonempty("version-baseband", var_all)
        self.exists_nonempty("version-bootloader", var_all)
        self.exists_nonempty("product", var_all)
        self.exists_yes_no("secure", var_all)
        self.exists_yes_no("unlocked", var_all)
        self.exists_validvals("off-mode-charge", var_all, ["0", "1"])
        self.assertIn("variant", var_all)
        voltage = self.exists_nonempty("battery-voltage", var_all)
        if voltage[-2:].lower() == "mv":
            voltage = voltage[:-2]
        try:
            voltnum = float(voltage)
        except ValueError:
            self.fail("battery-voltage (%s) is not a number" % (varname, voltage))
        self.exists_yes_no("battery-soc-ok", var_all)
        maxdl = self.exists_integer("max-download-size", var_all, 16)
        self.assertGreater(maxdl, 0)

        if "slot-count" in var_all:
            try:
                slotcount = int(var_all["slot-count"])
            except ValueError:
                self.fail("slot-count (%s) is not an integer" % var_all["slot-count"])
            if slotcount > 1:
                # test for A/B variables
                slots = [chr(slotnum+ord('a')) for slotnum in range(slotcount)]
                self.exists_validvals("current-slot", var_all, slots)

                # test for slot metadata
                for slot in slots:
                    self.exists_yes_no("slot-unbootable:"+slot, var_all)
                    self.exists_yes_no("slot-unbootable:"+slot, var_all)
                    self.exists_integer("slot-retry-count:"+slot, var_all)
            else:
                print "This does not appear to be an A/B device."

    def test_getvar_nonexistent(self):
        """Tests behaviour of nonexistent variables."""

        self.assertIsNone(self.fastboot.getvar("fhqwhgads"))

    def test_getvar(self):
        """Tests all variables separately"""

        self.get_exists_nonempty("version-baseband")
        self.get_exists_nonempty("version-bootloader")
        self.get_exists_nonempty("product")
        self.get_exists_yes_no("secure")
        self.get_exists_yes_no("unlocked")
        self.get_exists_validvals("off-mode-charge", ["0", "1"])
        self.get_exists("variant")
        voltage = self.get_exists_nonempty("battery-voltage")
        if voltage[-2:].lower() == "mv":
            voltage = voltage[:-2]
        try:
            voltnum = float(voltage)
        except ValueError:
            self.fail("battery-voltage (%s) is not a number" % voltage)
        self.get_exists_yes_no("battery-soc-ok")
        maxdl = self.get_exists_integer("max-download-size", 16)
        self.assertGreater(maxdl, 0)

        slotcount = 0
        try:
            slotcountString = self.fastboot.getvar("slot-count")
            if slotcountString != None:
                slotcount = int(slotcountString)
        except ValueError:
            self.fail("slot-count (%s) is not an integer" % slotcountString)
        if slotcount  > 1:
            # test for A/B variables
            slots = [chr(slotnum+ord('a')) for slotnum in range(slotcount)]
            self.get_exists_validvals("current-slot", slots)

            # test for slot metadata
            for slot in slots:
                self.get_exists_yes_no("slot-unbootable:"+slot)
                self.get_exists_yes_no("slot-successful:"+slot)
                self.get_exists_integer("slot-retry-count:"+slot)
        else:
            print "This does not appear to be an A/B device."

    def test_setactive(self):
        """Tests that A/B devices can switch to each slot, and the change persists over a reboot."""

        slotcount = 0
        try:
            val = self.fastboot.getvar("slot-count")
            if val != None:
                slotcount = int(val)
        except ValueError:
            self.fail("slot-count (%s) is not an integer" % val)
        except subprocess.CalledProcessError:
            print "Does not appear to be an A/B device."
        maxtries = 0
        if slotcount > 1:
            slots = [chr(slotnum+ord('a')) for slotnum in range(slotcount)]
            for slot in slots:
                self.fastboot.set_active(slot)
                self.assertEqual(slot, self.fastboot.getvar("current-slot"))
                self.assertEqual("no", self.fastboot.getvar("slot-unbootable:"+slot))
                self.assertEqual("no", self.fastboot.getvar("slot-successful:"+slot))
                retry = self.get_exists_integer("slot-retry-count:"+slot)
                if maxtries == 0:
                   maxtries = retry
                else:
                   self.assertEqual(maxtries, retry)
                self.fastboot.reboot(True)
                self.assertEqual(slot, self.fastboot.getvar("current-slot"))
                self.assertEqual("no", self.fastboot.getvar("slot-unbootable:"+slot))
                self.assertEqual("no", self.fastboot.getvar("slot-successful:"+slot))
                retry = self.get_exists_integer("slot-retry-count:"+slot)
                if maxtries == 0:
                   maxtries = retry
                else:
                   self.assertEqual(maxtries, retry)
        else:
            print "Does not appear to be an A/B device."

if __name__ == '__main__':
    unittest.main(verbosity=3)