#!/usr/bin/python3

# switcheroo-control 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>
# (C) 2020 Bastien Nocera <hadess@hadess.net>
#
# 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 dbus
import tempfile
import subprocess
import unittest
import time

try:
    import gi
    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)

try:
    gi.require_version('UMockdev', '1.0')
    from gi.repository import UMockdev
except ImportError:
    sys.stderr.write('Skipping tests, umockdev not available (https://github.com/martinpitt/umockdev)\n')
    sys.exit(0)

try:
    import dbusmock
except ImportError:
    sys.stderr.write('Skipping tests, python-dbusmock not available (http://pypi.python.org/pypi/python-dbusmock).\n')
    sys.exit(0)


SC = 'net.hadess.SwitcherooControl'
SC_PATH = '/net/hadess/SwitcherooControl'

class Tests(dbusmock.DBusTestCase):
    @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', 'switcheroo-control'), os.X_OK):
            cls.daemon_path = os.path.join(builddir, 'src', 'switcheroo-control')
            print('Testing binaries from local build tree (%s)' % cls.daemon_path)
            cls.local_daemon = True
        elif os.environ.get('UNDER_JHBUILD', False):
            jhbuild_prefix = os.environ['JHBUILD_PREFIX']
            cls.daemon_path = os.path.join(jhbuild_prefix, 'libexec', 'switcheroo-control')
            print('Testing binaries from JHBuild (%s)' % cls.daemon_path)
            cls.local_daemon = False
        else:
            cls.daemon_path = None
            with open('/usr/lib/systemd/system/switcheroo-control.service') as f:
                for line in f:
                    if line.startswith('ExecStart='):
                        cls.daemon_path = line.split('=', 1)[1].strip()
                        break
            assert cls.daemon_path, 'could not determine daemon path from systemd .service file'
            cls.local_daemon = False
            print('Testing installed system binary (%s)' % cls.daemon_path)

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

        # set up a fake system D-BUS
        cls.test_bus = Gio.TestDBus.new(Gio.TestDBusFlags.NONE)
        cls.test_bus.up()
        try:
            del os.environ['DBUS_SESSION_BUS_ADDRESS']
        except KeyError:
            pass
        os.environ['DBUS_SYSTEM_BUS_ADDRESS'] = cls.test_bus.get_bus_address()

        cls.dbus = Gio.bus_get_sync(Gio.BusType.SYSTEM, None)
        cls.dbus_con = cls.get_dbus(True)

    @classmethod
    def tearDownClass(cls):
        cls.test_bus.down()
        dbusmock.DBusTestCase.tearDownClass()

    def setUp(self):
        '''Set up a local umockdev testbed.

        The testbed is initially empty.
        '''
        self.testbed = UMockdev.Testbed.new()

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

    def tearDown(self):
        del self.testbed
        self.stop_daemon()

        # on failures, print daemon log
        errors = [x[1] for x in self._outcome.errors if x[1]]
        if errors and self.log:
            with open(self.log.name) as f:
                sys.stderr.write('\n-------------- daemon log: ----------------\n')
                sys.stderr.write(f.read())
                sys.stderr.write('------------------------------\n')

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

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

        When done, this sets self.proxy as the Gio.DBusProxy for switcheroo-control.
        '''
        env = os.environ.copy()
        env['G_DEBUG'] = 'fatal-criticals'
        env['G_MESSAGES_DEBUG'] = 'all'
        # note: Python doesn't propagate the setenv from Testbed.new(), so we
        # have to do that ourselves
        env['UMOCKDEV_DIR'] = self.testbed.get_root_dir()
        self.log = tempfile.NamedTemporaryFile()
        if os.getenv('VALGRIND') != None:
            if self.local_daemon:
                daemon_path = ['libtool', '--mode=execute', 'valgrind', self.daemon_path, '-v']
            else:
                daemon_path = ['valgrind', self.daemon_path, '-v']
        else:
            daemon_path = [self.daemon_path, '-v']

        self.daemon = subprocess.Popen(daemon_path,
                                       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('HasDualGpu')
                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, SC,
            SC_PATH, SC, 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, SC,
            SC_PATH, 'org.freedesktop.DBus.Properties', None)
        return proxy.Get('(ss)', SC, name)

    def have_text_in_log(self, text):
        return self.count_text_in_log(text) > 0

    def count_text_in_log(self, text):
        with open(self.log.name) as f:
            return f.read().count(text)

    def assertEventually(self, condition, message=None, timeout=50):
        '''Assert that condition function eventually returns True.

        Timeout is in deciseconds, defaulting to 50 (5 seconds). message is
        printed on failure.
        '''
        while timeout >= 0:
            context = GLib.MainContext.default()
            while context.iteration(False):
                pass
            if condition():
                break
            timeout -= 1
            time.sleep(0.1)
        else:
            self.fail(message or 'timed out waiting for ' + str(condition))

    def add_intel_gpu(self):
        parent = self.testbed.add_device('pci', 'i915 VGA controller', None,
                [ 'boot_vga', '1' ],
                [ 'DRIVER', 'i915',
                  'PCI_CLASS', '30000',
                  'PCI_ID', '8086:5917',
                  'PCI_SUBSYS_ID', '1043:1A00'
                  'PCI_SLOT_NAME', '0000:00:02.0'
                  'MODALIAS', 'pci:v00008086d00005917sv00001043sd00001A00bc03sc00i00',
                  'ID_PCI_CLASS_FROM_DATABASE', 'Display controller',
                  'ID_PCI_SUBCLASS_FROM_DATABASE', 'VGA compatible controller',
                  'ID_PCI_INTERFACE_FROM_DATABASE', 'VGA controller',
                  'ID_VENDOR_FROM_DATABASE', 'Intel Corporation',
                  'ID_MODEL_FROM_DATABASE', 'UHD Graphics 620',
                  'SWITCHEROO_CONTROL_PRODUCT_NAME', 'UHD Graphics 620 (Kabylake GT2)',
                  'SWITCHEROO_CONTROL_VENDOR_NAME', 'Intel(R)',
                  'FWUPD_GUID', '0x8086:0x5917' ]
                )

        self.testbed.add_device('drm', 'dri/card0', parent,
                [],
                [ 'DEVNAME', '/dev/dri/card0',
                  'ID_PATH', 'pci-0000:00:02.0',
                  'ID_PATH_TAG', 'pci-0000_00_02_0' ]
                )

        self.testbed.add_device('drm', 'dri/renderD128', parent,
                [],
                [ 'DEVNAME', '/dev/dri/renderD128',
                  'ID_PATH', 'pci-0000:00:02.0',
                  'ID_PATH_TAG', 'pci-0000_00_02_0' ]
                )

    def add_nouveau_gpu(self):
        parent = self.testbed.add_device('pci', 'NVidia VGA controller', None,
                [ 'boot_vga', '0' ],
                [ 'DRIVER', 'nouveau',
                  'PCI_CLASS', '30200',
                  'PCI_ID', '10DE:134E',
                  'PCI_SUBSYS_ID', '1043:143E'
                  'PCI_SLOT_NAME', '0000:01:00.0'
                  'MODALIAS', 'pci:v000010DEd0000134Esv00001043sd0000143Ebc03sc02i00',
                  'ID_PCI_CLASS_FROM_DATABASE', 'Display controller',
                  'ID_PCI_SUBCLASS_FROM_DATABASE', '3D controller',
                  'ID_PCI_INTERFACE_FROM_DATABASE', 'NVIDIA Corporation',
                  'ID_MODEL_FROM_DATABASE', 'GM108M [GeForce 930MX]',
                  'FWUPD_GUID', '0x10de:0x134e' ]
                )

        self.testbed.add_device('drm', 'dri/card1', parent,
                [],
                [ 'DEVNAME', '/dev/dri/card1',
                  'ID_PATH', 'pci-0000:01:00.0',
                  'ID_PATH_TAG', 'pci-0000_01_00_0' ]
                )

        self.testbed.add_device('drm', 'dri/renderD129', parent,
                [],
                [ 'DEVNAME', '/dev/dri/renderD129',
                  'ID_PATH', 'pci-0000:01:00.0',
                  'ID_PATH_TAG', 'pci-0000_01_00_0' ]
                )

    #
    # Actual test cases
    #

    def test_single_device(self):
        '''single device'''

        self.add_intel_gpu()

        self.start_daemon()
        self.assertEqual(self.get_dbus_property('HasDualGpu'), False)
        self.assertEqual(self.get_dbus_property('NumGPUs'), 1)

        gpus = self.get_dbus_property('GPUs')
        self.assertEqual(len(gpus), 1)
        self.assertEqual(gpus[0]['Name'], 'Intel® UHD Graphics 620 (Kabylake GT2)')
        sc_env = gpus[0]['Environment']
        self.assertEqual(len(sc_env), 2)
        self.assertEqual(sc_env[0], 'DRI_PRIME')
        self.assertEqual(sc_env[1], 'pci-0000_00_02_0')
        self.assertEqual(gpus[0]['Default'], True)

        # process = subprocess.Popen(['gdbus', 'introspect', '--system', '--dest', 'net.hadess.SwitcherooControl', '--object-path', '/net/hadess/SwitcherooControl'])
        # print (self.get_dbus_property('GPUs'))

        self.stop_daemon()

    def test_dual_open_source(self):
        '''dual open source devices'''

        self.add_intel_gpu()
        self.add_nouveau_gpu()

        self.start_daemon()
        self.assertEqual(self.get_dbus_property('HasDualGpu'), True)
        self.assertEqual(self.get_dbus_property('NumGPUs'), 2)

        gpus = self.get_dbus_property('GPUs')
        self.assertEqual(len(gpus), 2)

        gpu = gpus[0]
        self.assertEqual(gpu['Name'], 'GM108M [GeForce 930MX]')
        sc_env = gpu['Environment']
        self.assertEqual(len(sc_env), 2)
        self.assertEqual(sc_env[0], 'DRI_PRIME')
        self.assertEqual(sc_env[1], 'pci-0000_01_00_0')
        self.assertEqual(gpu['Default'], False)

        gpu = gpus[1]
        self.assertEqual(gpu['Name'], 'Intel® UHD Graphics 620 (Kabylake GT2)')
        sc_env = gpu['Environment']
        self.assertEqual(len(sc_env), 2)
        self.assertEqual(sc_env[0], 'DRI_PRIME')
        self.assertEqual(sc_env[1], 'pci-0000_00_02_0')
        self.assertEqual(gpu['Default'], True)

        # process = subprocess.Popen(['gdbus', 'introspect', '--system', '--dest', 'net.hadess.SwitcherooControl', '--object-path', '/net/hadess/SwitcherooControl'])

        self.stop_daemon()

    def test_dual_open_source(self):
        '''dual open source devices'''

        self.add_intel_gpu()
        self.add_nouveau_gpu()

        self.testbed.add_device('drm', 'ttm', None,
                [],
                [ 'DEVPATH', '/devices/virtual/drm/ttm',
                  'DEVTYPE', 'ttm' ]
                )

        self.start_daemon()
        self.assertEqual(self.get_dbus_property('HasDualGpu'), True)
        self.assertEqual(self.get_dbus_property('NumGPUs'), 2)

        gpus = self.get_dbus_property('GPUs')
        self.assertEqual(len(gpus), 2)

        gpu = gpus[0]
        self.assertEqual(gpu['Name'], 'GM108M [GeForce 930MX]')
        self.assertEqual(gpu['Default'], False)

        gpu = gpus[1]
        self.assertEqual(gpu['Name'], 'Intel® UHD Graphics 620 (Kabylake GT2)')
        self.assertEqual(gpu['Default'], True)

        # process = subprocess.Popen(['gdbus', 'introspect', '--system', '--dest', 'net.hadess.SwitcherooControl', '--object-path', '/net/hadess/SwitcherooControl'])

        self.stop_daemon()

    def test_dual_hotplug(self):
        '''dual open source devices'''

        self.add_intel_gpu()

        self.start_daemon()
        self.assertEqual(self.get_dbus_property('HasDualGpu'), False)
        self.assertEqual(self.get_dbus_property('NumGPUs'), 1)

        self.add_nouveau_gpu()

        self.assertEqual(self.get_dbus_property('HasDualGpu'), True)
        self.assertEqual(self.get_dbus_property('NumGPUs'), 2)

        # process = subprocess.Popen(['gdbus', 'introspect', '--system', '--dest', 'net.hadess.SwitcherooControl', '--object-path', '/net/hadess/SwitcherooControl'])

        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__':
    # run ourselves under umockdev
    if 'umockdev' not in os.environ.get('LD_PRELOAD', ''):
        os.execvp('umockdev-wrapper', ['umockdev-wrapper'] + sys.argv)

    unittest.main()
