#!/usr/bin/python3

# upower integration test suite
#
# Run in built tree to test local built binaries, or from anywhere else to test
# system installed binaries.
#
# Copyright: (C) 2011 Martin Pitt <martin.pitt@ubuntu.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

import os
import sys
import tempfile
import subprocess
import unittest
import shutil
import time

try:
    from gi.repository import GLib
    from gi.repository import Gio
except ImportError as e:
    sys.stderr.write('Skipping tests, PyGobject not available for Python 3, or missing GI typelibs: %s\n' % str(e))
    sys.exit(0)

UP = 'org.freedesktop.UPower'

(UP_DEVICE_STATE_UNKNOWN, 
 UP_DEVICE_STATE_CHARGING,
 UP_DEVICE_STATE_DISCHARGING,
 UP_DEVICE_STATE_EMPTY,
 UP_DEVICE_STATE_FULLY_CHARGED) = (0, 1, 2, 3, 4)

class Tests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # run from local build tree if we are in one, otherwise use system instance
        builddir = os.getenv('top_builddir', '.')
        if os.access (os.path.join(builddir, 'src', 'upowerd'), os.X_OK):
            daemon_path = os.path.join(builddir, 'src', 'upowerd')
            print('Testing binaries from local build tree')
        else:
            print('Testing installed system binaries')
            daemon_path = None
            with open('/usr/share/dbus-1/system-services/org.freedesktop.UPower.service') as f:
                for line in f:
                    if line.startswith('Exec='):
                        daemon_path = line.split('=', 1)[1].strip()
                        break
            assert daemon_path, 'could not determine daemon path from D-BUS .service file'

        # fail on CRITICALs on client side
        GLib.log_set_always_fatal(GLib.LogLevelFlags.LEVEL_WARNING|
                                  GLib.LogLevelFlags.LEVEL_ERROR|
                                  GLib.LogLevelFlags.LEVEL_CRITICAL)

        # if we are root, test the real thing on the system bus, otherwise
        # start on the session bus
        if os.geteuid() == 0:
            # kill running daemons
            subprocess.call(['killall', 'upowerd'])

            cls.dbus = Gio.bus_get_sync(Gio.BusType.SYSTEM, None)
            cls.daemon_argv = [daemon_path]
            print('Testing as root on the system bus')
        else:
            cls.dbus = Gio.bus_get_sync(Gio.BusType.SESSION, None)
            cls.daemon_argv = [daemon_path, '--test']

            # use dbus-lauch if possible
            print('Testing as user on the session bus')


    def setUp(self):
        '''Set up a local sysfs tree and determine paths.
        
        The sysfs tree is empty initially and needs to be populated with
        @add_device.
        '''
        self.sysfs = tempfile.mkdtemp()

        self.proxy = None 
        self.log = None
        self.daemon = None

    def tearDown(self):
        self.stop_daemon()
        shutil.rmtree(self.sysfs)

        # on failures, print daemon log
        if not self._outcomeForDoCleanups.success:
            with open(self.log.name) as f:
                sys.stderr.write('\n-------------- daemon log: ----------------\n')
                sys.stderr.write(f.read())
                sys.stderr.write('------------------------------\n')

    #
    # Methods for fake sysfs
    #

    def add_device(self, subsystem, name, attributes, properties=None):
        '''Add a new device to the local sysfs tree.
        
        Return the device path.
        '''
        dev_dir = os.path.join(self.sysfs, 'devices', name)
        if not os.path.isdir(dev_dir):
            os.makedirs(dev_dir)
        class_dir = os.path.join(self.sysfs, 'class', subsystem)
        if not os.path.isdir(class_dir):
            os.makedirs(class_dir)

        os.symlink(os.path.join('..', '..', 'devices', name), os.path.join(class_dir, name))
        os.symlink(os.path.join('..', '..', 'class', subsystem), os.path.join(dev_dir, 'subsystem'))

        attributes['uevent'] = self._props_to_str(properties)

        for a, v in attributes.items():
            self.set_sys_attribute(dev_dir, a, v)

        return dev_dir

    def get_sys_attribute(self, devpath, name):
        '''Get device attribute from the local sysfs tree.'''

        with open(os.path.join(devpath, name), 'r') as f:
            return f.read()

    def set_sys_attribute(self, devpath, name, value):
        '''Set device attribute in the local sysfs tree.'''

        with open(os.path.join(devpath, name), 'w') as f:
            f.write(value)

    def set_sys_property(self, devpath, name, value):
        '''Set device udev property in the local sysfs tree.'''

        prop_str = self.get_sys_attribute(devpath, 'uevent')
        props = {}
        for l in prop_str.splitlines():
            (k, v) = l.split('=')
            props[k] = v.rstrip()

        props[name] = value

        self.set_sys_attribute(devpath, 'uevent', self._props_to_str(props))

    #
    # Daemon control and D-BUS I/O
    #

    def start_daemon(self):
        '''Start daemon and create DBus proxy.

        Do this after adding the devices you want to test with. At the moment
        this only works with coldplugging, as we do not currently have a way to
        inject simulated uevents.

        When done, this sets self.proxy a the Gio.DBusProxy for upowerd.
        '''
        env = os.environ.copy()
        env['G_DEBUG'] = 'fatal-criticals'
        env['SYSFS_PATH'] = self.sysfs
        self.log = tempfile.NamedTemporaryFile()
        self.daemon = subprocess.Popen(self.daemon_argv,
                env=env, stdout=self.log, stderr=subprocess.STDOUT)

        # wait until the daemon gets online
        timeout = 100
        while timeout > 0:
            time.sleep(0.1)
            timeout -= 1
            try:
                self.get_dbus_property('DaemonVersion')
                break
            except GLib.GError:
                pass
        else:
            self.fail('daemon did not start in 10 seconds')

        self.proxy = Gio.DBusProxy.new_sync(self.dbus,
                Gio.DBusProxyFlags.DO_NOT_AUTO_START, None,
                UP, '/org/freedesktop/UPower', UP, None)

        self.assertEqual(self.daemon.poll(), None, 'daemon crashed')

    def stop_daemon(self):
        '''Stop the daemon if it is running.'''

        if self.daemon:
            try:
                self.daemon.kill()
            except OSError:
                pass
            self.daemon.wait()
        self.daemon = None
        self.proxy = None

    def get_dbus_property(self, name):
        '''Get property value from daemon D-Bus interface.'''

        proxy = Gio.DBusProxy.new_sync(self.dbus,
                Gio.DBusProxyFlags.DO_NOT_AUTO_START, None,
                UP, '/org/freedesktop/UPower', 
                'org.freedesktop.DBus.Properties', None)
        return proxy.Get('(ss)', UP, name)

    def get_dbus_dev_property(self, device, name):
        '''Get property value from an upower device D-Bus path.'''

        proxy = Gio.DBusProxy.new_sync(self.dbus,
                Gio.DBusProxyFlags.DO_NOT_AUTO_START, None,
                UP, device, 'org.freedesktop.DBus.Properties', None)
        return proxy.Get('(ss)', UP + '.Device', name)

    #
    # Actual test cases
    #

    def test_daemon_version(self):
        '''DaemonVersion property'''

        self.start_daemon()
        self.assertEqual(self.proxy.EnumerateDevices(), [])
        self.assertRegex(self.get_dbus_property('DaemonVersion'), '^[0-9.]+$')

        # without any devices we should assume AC
        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)

    def test_battery_ac(self):
        '''battery properties with and without AC'''

        # without any devices we should assume AC
        self.start_daemon()
        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

        # online AC
        ac = self.add_device('power_supply', 'AC', 
                {'type': 'Mains', 'online': '1' })

        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 1)
        ac_up = devs[0]
        self.assertTrue('line_power_AC' in ac_up)
        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.assertEqual(self.get_dbus_dev_property(ac_up, 'PowerSupply'), True)
        self.assertEqual(self.get_dbus_dev_property(ac_up, 'Online'), True)
        self.assertEqual(self.get_dbus_dev_property(ac_up, 'NativePath'), ac)
        self.stop_daemon()

        # offline AC
        self.set_sys_attribute(ac, 'online', '0')
        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 1)
        # we don't have any known online power device now, but still no battery
        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.assertEqual(self.get_dbus_dev_property(ac_up, 'Online'), False)
        self.stop_daemon()

        # offline AC + discharging battery
        bat0 = self.add_device('power_supply', 'BAT0', 
                {'type': 'Battery', 
                 'present': '1', 
                 'status': 'Discharging',
                 'energy_full': '60000000',
                 'energy_full_design': '80000000',
                 'energy_now': '48000000',
                 'voltage_now': '12000000'})

        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 2)
        if devs[0] == ac_up:
            bat0_up = devs[1]
        else:
            bat0_up = devs[0]
        # we don't have any known online power device now, but still no battery
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'IsPresent'), True)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'State'), UP_DEVICE_STATE_DISCHARGING)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'Percentage'), 80.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'Energy'), 48.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'EnergyFull'), 60.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'EnergyFullDesign'), 80.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'Voltage'), 12.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'NativePath'), bat0)
        self.stop_daemon()

        # offline AC + discharging low battery
        self.set_sys_attribute(bat0, 'energy_now', '1500000')
        self.start_daemon()
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), True)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'IsPresent'), True)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'State'), UP_DEVICE_STATE_DISCHARGING)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'Percentage'), 2.5)
        self.stop_daemon()

        # now connect AC again
        self.set_sys_attribute(ac, 'online', '1')
        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 2)
        # we don't have any known online power device now, but still no battery
        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.assertEqual(self.get_dbus_dev_property(ac_up, 'Online'), True)
        self.stop_daemon()

    def test_multiple_batteries(self):
        '''Multiple batteries'''

        # one well charged, one low
        bat0 = self.add_device('power_supply', 'BAT0', 
                {'type': 'Battery', 
                 'present': '1', 
                 'status': 'Discharging',
                 'energy_full': '60000000',
                 'energy_full_design': '80000000',
                 'energy_now': '48000000',
                 'voltage_now': '12000000'})

        bat1 = self.add_device('power_supply', 'BAT1', 
                {'type': 'Battery', 
                 'present': '1', 
                 'status': 'Discharging',
                 'energy_full': '60000000',
                 'energy_full_design': '80000000',
                 'energy_now': '1500000',
                 'voltage_now': '12000000'})

        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 2)

        # as we have one which is well-charged, the summary state is "not low
        # battery"
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

        # now set both to low
        self.set_sys_attribute(bat0, 'energy_now', '1500000')
        self.start_daemon()
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), True)
        self.stop_daemon()

    def test_unknown_battery_status(self):
        '''Unknown battery charge status'''

        bat0 = self.add_device('power_supply', 'BAT0', 
                {'type': 'Battery', 
                 'present': '1', 
                 'status': 'unknown',
                 'energy_full': '60000000',
                 'energy_full_design': '80000000',
                 'energy_now': '48000000',
                 'voltage_now': '12000000'})

        # with no other power sources, the OnBattery value here is really
        # arbitrary, so don't test it. The only thing we know for sure is that
        # we aren't on low battery
        self.start_daemon()
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

        # However, if we have an AC, we can infer
        ac = self.add_device('power_supply', 'AC', 
                {'type': 'Mains', 'online': '0' })
        self.start_daemon()
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

        self.set_sys_attribute(ac, 'online', '1')
        self.start_daemon()
        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

    def test_battery_charge(self):
        '''battery which reports charge instead of energy
        
        energy_* is in uWh, while charge_* is in uAh.
        '''
        bat0 = self.add_device('power_supply', 'BAT0', 
                {'type': 'Battery', 
                 'present': '1', 
                 'status': 'Discharging',
                 'charge_full': '10500000',
                 'charge_full_design': '11000000',
                 'charge_now': '7875000',
                 'current_now': '787000',
                 'voltage_now': '12000000'})

        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 1)
        bat0_up = devs[0]

        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'IsPresent'), True)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'State'), UP_DEVICE_STATE_DISCHARGING)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'Percentage'), 75.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'Energy'), 94.5)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'EnergyFull'), 126.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'EnergyFullDesign'), 132.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'Voltage'), 12.0)
        self.assertEqual(self.get_dbus_dev_property(bat0_up, 'NativePath'), bat0)
        self.stop_daemon()

    def test_ups_ac(self):
        '''UPS properties with and without AC'''

        # add a charging UPS
        ups0 = self.add_device('usb', 'hiddev0', {}, 
                {'DEVNAME': 'null', 'UPOWER_VENDOR': 'APC',
                 'UPOWER_BATTERY_TYPE': 'ups',
                 'UPOWER_FAKE_DEVICE': '1',
                 'UPOWER_FAKE_HID_CHARGING': '1',
                 'UPOWER_FAKE_HID_PERCENTAGE': '70'})

        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 1)
        ups0_up = devs[0]

        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'IsPresent'), True)
        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'Percentage'), 70.0)
        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'State'), UP_DEVICE_STATE_CHARGING)

        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

        # now switch to discharging UPS
        self.set_sys_property(ups0, 'UPOWER_FAKE_HID_CHARGING', '0')

        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 1)
        self.assertEqual(devs[0], ups0_up)

        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'IsPresent'), True)
        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'Percentage'), 70.0)
        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'State'), UP_DEVICE_STATE_DISCHARGING)
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

        # low UPS charge
        self.set_sys_property(ups0, 'UPOWER_FAKE_HID_PERCENTAGE', '2')
        self.start_daemon()
        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'Percentage'), 2.0)
        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'State'), UP_DEVICE_STATE_DISCHARGING)
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), True)
        self.stop_daemon()

        # now add an offline AC, should still be on battery
        ac = self.add_device('power_supply', 'AC', 
                {'type': 'Mains', 'online': '0' })
        self.start_daemon()
        devs = self.proxy.EnumerateDevices()
        self.assertEqual(len(devs), 2)
        if devs[0] == ups0_up:
            ac_up = devs[1]
        else:
            ac_up = devs[0]

        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'Percentage'), 2.0)
        self.assertEqual(self.get_dbus_dev_property(ups0_up, 'State'), UP_DEVICE_STATE_DISCHARGING)
        self.assertEqual(self.get_dbus_property('OnBattery'), True)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), True)
        self.stop_daemon()

        # now plug in the AC, should switch to OnBattery=False
        self.set_sys_attribute(ac, 'online', '1')
        self.start_daemon()
        self.assertEqual(self.get_dbus_property('OnBattery'), False)
        self.assertEqual(self.get_dbus_property('OnLowBattery'), False)
        self.stop_daemon()

    #
    # Helper methods
    #

    @classmethod
    def _props_to_str(cls, properties):
        '''Convert a properties dictionary to uevent text representation.'''

        prop_str = ''
        if properties:
            for k, v in properties.items():
                prop_str += '%s=%s\n' % (k, v)
        return prop_str

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