#!/usr/bin/python3
# -*- Mode: Python; indent-tabs-mode: nil; tab-width: 4; coding: utf-8 -*-

import apt_pkg
import mock
import os
import unittest
import tempfile
import json

from DistUpgrade.DistUpgradeQuirks import DistUpgradeQuirks

CURDIR = os.path.dirname(os.path.abspath(__file__))


class MockController(object):
    def __init__(self):
        self._view = None


class MockConfig(object):
    pass


class MockPopenSnap():
    def __init__(self, cmd, universal_newlines=True,
                 stdout=None):
        self.command = cmd

    def communicate(self):
        if self.command[1] == "list":
            return []
        snap_name = self.command[2]
        if snap_name == 'gnome-logs':
            # Package to refresh
            return ["""
name:      test-snap
summary:   Test Snap
publisher: Canonical
license:   unset
description: Some description
commands:
  - gnome-calculator
snap-id:      1234
tracking:     stable/ubuntu-19.04
refresh-date: 2019-04-11
channels:
  stable:    3.32.1  2019-04-10 (406) 4MB -
  candidate: 3.32.2  2019-06-26 (433) 4MB -
  beta:      3.33.89 2019-08-06 (459) 4MB -
  edge:      3.33.90 2019-08-06 (460) 4MB -
installed:   3.32.1             (406) 4MB -
"""]
        elif "gnome-characters" in snap_name:
            # Package installed but not tracking the release channel
            return ["""
name:      test-snap
summary:   Test Snap
publisher: Canonical
license:   unset
description: Some description
commands:
  - gnome-characters
snap-id:      1234
refresh-date: 2019-04-11
channels:
  stable:    3.32.1  2019-04-10 (406) 4MB -
  candidate: 3.32.2  2019-06-26 (433) 4MB -
  beta:      3.33.89 2019-08-06 (459) 4MB -
  edge:      3.33.90 2019-08-06 (460) 4MB -
installed:   3.32.1             (406) 4MB -
"""]
        elif "gtk-common-themes" in snap_name:
            # Package installed but missing/invalid snap id
            return ["""
name:      test-snap
summary:   Test Snap
publisher: Canonical
license:   unset
description: Some description
commands:
  - gtk-common-themes
refresh-date: 2019-04-11
channels:
  stable:    3.32.1  2019-04-10 (406) 4MB -
  candidate: 3.32.2  2019-06-26 (433) 4MB -
  beta:      3.33.89 2019-08-06 (459) 4MB -
  edge:      3.33.90 2019-08-06 (460) 4MB -
"""]
        else:
            return ["""
name:      test-snap
summary:   Test Snap
publisher: Canonical
license:   unset
description: Some description
commands:
  - gnome-calculator
snap-id:      1234
refresh-date: 2019-04-11
channels:
  stable:    3.32.1  2019-04-10 (406) 4MB -
  candidate: 3.32.2  2019-06-26 (433) 4MB -
  beta:      3.33.89 2019-08-06 (459) 4MB -
  edge:      3.33.90 2019-08-06 (460) 4MB -
"""]


def mock_urlopen_snap(req):
    result = """{{
  "error-list": [],
  "results": [
    {{
      "effective-channel": "stable",
      "instance-key": "test",
      "name": "{name}",
      "released-at": "2019-04-10T18:54:15.717357+00:00",
      "result": "download",
      "snap": {{
        "created-at": "2019-04-09T17:09:29.941588+00:00",
        "download": {{
          "deltas": [],
          "size": {size},
          "url": "SNAPURL"
        }},
        "license": "GPL-3.0+",
        "name": "{name}",
        "prices": {{ }},
        "publisher": {{
          "display-name": "Canonical",
          "id": "canonical",
          "username": "canonical",
          "validation": "verified"
        }},
        "revision": 406,
        "snap-id": "{snap_id}",
        "summary": "GNOME Calculator",
        "title": "GNOME Calculator",
        "type": "app",
        "version": "3.32.1"
      }},
      "snap-id": "{snap_id}"
    }}
  ]
}}
"""
    test_snaps = {
        '1': ("gnome-calculator", 4218880),
        '2': ("test-snap", 2000000)
    }
    json_data = json.loads(req.data)
    snap_id = json_data['actions'][0]['snap-id']
    name = test_snaps[snap_id][0]
    size = test_snaps[snap_id][1]
    response_mock = mock.Mock()
    response_mock.read.return_value = result.format(
        name=name, snap_id=snap_id, size=size)
    return response_mock


def make_mock_pkg(name, is_installed, candidate_rec=""):
    mock_pkg = mock.Mock()
    mock_pkg.name = name
    mock_pkg.is_installed = is_installed
    mock_pkg.marked_install = False
    if candidate_rec:
        mock_pkg.candidate = mock.Mock()
        mock_pkg.candidate.record = candidate_rec
    return mock_pkg


class TestQuirks(unittest.TestCase):

    orig_recommends = ''
    orig_status = ''

    def setUp(self):
        self.orig_recommends = apt_pkg.config.get("APT::Install-Recommends")
        self.orig_status = apt_pkg.config.get("Dir::state::status")

    def tearDown(self):
        apt_pkg.config.set("APT::Install-Recommends", self.orig_recommends)
        apt_pkg.config.set("Dir::state::status", self.orig_status)

    def test_enable_recommends_during_upgrade(self):
        controller = mock.Mock()

        config = mock.Mock()
        q = DistUpgradeQuirks(controller, config)
        # server mode
        apt_pkg.config.set("APT::Install-Recommends", "0")
        controller.serverMode = True
        self.assertFalse(apt_pkg.config.find_b("APT::Install-Recommends"))
        q.ensure_recommends_are_installed_on_desktops()
        self.assertFalse(apt_pkg.config.find_b("APT::Install-Recommends"))
        # desktop mode
        apt_pkg.config.set("APT::Install-Recommends", "0")
        controller.serverMode = False
        self.assertFalse(apt_pkg.config.find_b("APT::Install-Recommends"))
        q.ensure_recommends_are_installed_on_desktops()
        self.assertTrue(apt_pkg.config.find_b("APT::Install-Recommends"))

    def test_screensaver_poke(self):
        # fake nothing is installed
        empty_status = tempfile.NamedTemporaryFile()
        apt_pkg.config.set("Dir::state::status", empty_status.name)

        # create quirks class
        controller = mock.Mock()
        config = mock.Mock()
        quirks = DistUpgradeQuirks(controller, config)
        quirks._pokeScreensaver()
        res = quirks._stopPokeScreensaver()
        res  # pyflakes

    def test_get_linux_metapackage(self):
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())
        mock_cache = set([
            make_mock_pkg(
                name="linux-image-3.19-24-generic",
                is_installed=True,
                candidate_rec={"Source": "linux"},
            ),
        ])
        pkgname = q._get_linux_metapackage(mock_cache, headers=False)
        self.assertEqual(pkgname, "linux-generic")

    def test_get_lpae_linux_metapackage(self):
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())
        mock_cache = set([
            make_mock_pkg(
                name="linux-image-4.2.0-16-generic-lpae",
                is_installed=True,
                candidate_rec={"Source": "linux"},
            ),
        ])
        pkgname = q._get_linux_metapackage(mock_cache, headers=False)
        self.assertEqual(pkgname, "linux-generic-lpae")

    def test_get_lowlatency_linux_metapackage(self):
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())
        mock_cache = set([
            make_mock_pkg(
                name="linux-image-4.2.0-16-lowlatency",
                is_installed=True,
                candidate_rec={"Source": "linux"},
            ),
        ])
        pkgname = q._get_linux_metapackage(mock_cache, headers=False)
        self.assertEqual(pkgname, "linux-lowlatency")

    def test_get_lts_linux_metapackage(self):
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())
        mock_cache = set([
            make_mock_pkg(
                name="linux-image-3.13.0-24-generic",
                is_installed=True,
                candidate_rec={"Source": "linux-lts-quantal"},
            ),
        ])
        pkgname = q._get_linux_metapackage(mock_cache, headers=False)
        self.assertEqual(pkgname, "linux-generic-lts-quantal")

    def test_ros_installed_warning(self):
        ros_packages = (
            "ros-melodic-catkin",
            "ros-noetic-rosboost-cfg",
            "ros-foxy-rosclean",
            "ros-kinetic-ros-environment",
            "ros-dashing-ros-workspace")
        for package_name in ros_packages:
            mock_controller = mock.Mock()
            mock_question = mock_controller._view.askYesNoQuestion
            mock_question.return_value = True

            q = DistUpgradeQuirks(mock_controller, mock.Mock())
            mock_cache = set([
                make_mock_pkg(
                    name=package_name,
                    is_installed=True,
                ),
            ])
            q._test_and_warn_if_ros_installed(mock_cache)
            mock_question.assert_called_once_with(mock.ANY, mock.ANY)
            self.assertFalse(len(mock_controller.abort.mock_calls))

            mock_controller.reset_mock()
            mock_question.reset_mock()
            mock_question.return_value = False

            mock_cache = set([
                make_mock_pkg(
                    name=package_name,
                    is_installed=True,
                ),
            ])
            q._test_and_warn_if_ros_installed(mock_cache)
            mock_question.assert_called_once_with(mock.ANY, mock.ANY)
            mock_controller.abort.assert_called_once_with()

    def test_ros_not_installed_no_warning(self):
        mock_controller = mock.Mock()
        mock_question = mock_controller._view.askYesNoQuestion
        mock_question.return_value = False

        q = DistUpgradeQuirks(mock_controller, mock.Mock())
        mock_cache = set([
            make_mock_pkg(
                name="ros-melodic-catkin",
                is_installed=False,
            ),
            make_mock_pkg(
                name="ros-noetic-rosboost-cfg",
                is_installed=False,
            ),
            make_mock_pkg(
                name="ros-foxy-rosclean",
                is_installed=False,
            ),
            make_mock_pkg(
                name="ros-kinetic-ros-environment",
                is_installed=False,
            ),
            make_mock_pkg(
                name="ros-dashing-ros-workspace",
                is_installed=False,
            ),
        ])
        q._test_and_warn_if_ros_installed(mock_cache)
        self.assertFalse(len(mock_question.mock_calls))
        self.assertFalse(len(mock_controller.abort.mock_calls))

    def test_replace_fkms_overlay_no_config(self):
        with tempfile.TemporaryDirectory() as boot_dir:
            mock_controller = mock.Mock()

            q = DistUpgradeQuirks(mock_controller, mock.Mock())

            q._replace_fkms_overlay(boot_dir)
            self.assertFalse(os.path.exists(os.path.join(
                boot_dir, 'config.txt.distUpgrade')))

    def test_replace_fkms_overlay_no_changes(self):
        with tempfile.TemporaryDirectory() as boot_dir:
            demo_config = """\
# This is a demo boot config
[pi4]
max_framebuffers=2
[all]
arm_64bit=1
kernel=vmlinuz
initramfs initrd.img followkernel"""
            with open(os.path.join(boot_dir, 'config.txt'), 'w') as f:
                f.write(demo_config)

            mock_controller = mock.Mock()

            q = DistUpgradeQuirks(mock_controller, mock.Mock())

            q._replace_fkms_overlay(boot_dir)
            self.assertFalse(os.path.exists(os.path.join(
                boot_dir, 'config.txt.distUpgrade')))
            with open(os.path.join(boot_dir, 'config.txt')) as f:
                self.assertTrue(f.read() == demo_config)

    def test_replace_fkms_overlay_with_changes(self):
        with tempfile.TemporaryDirectory() as boot_dir:
            demo_config = """\
# This is a demo boot config
[pi4]
max_framebuffers=2
[all]
arm_64bit=1
kernel=vmlinuz
initramfs initrd.img followkernel
dtoverlay=vc4-fkms-v3d,cma-256
start_x=1
gpu_mem=256
"""
            expected_config = """\
# This is a demo boot config
[pi4]
max_framebuffers=2
[all]
arm_64bit=1
kernel=vmlinuz
initramfs initrd.img followkernel
# changed by do-release-upgrade (LP: #1923673)
#dtoverlay=vc4-fkms-v3d,cma-256
dtoverlay=vc4-kms-v3d,cma-256
# disabled by do-release-upgrade (LP: #1923673)
#start_x=1
# disabled by do-release-upgrade (LP: #1923673)
#gpu_mem=256
"""
            with open(os.path.join(boot_dir, 'config.txt'), 'w') as f:
                f.write(demo_config)

            mock_controller = mock.Mock()

            q = DistUpgradeQuirks(mock_controller, mock.Mock())

            q._replace_fkms_overlay(boot_dir)
            self.assertTrue(os.path.exists(os.path.join(
                boot_dir, 'config.txt.distUpgrade')))
            with open(os.path.join(boot_dir, 'config.txt')) as f:
                self.assertTrue(f.read() == expected_config)

    def test_add_kms_overlay_no_config(self):
        with tempfile.TemporaryDirectory() as boot_dir:
            mock_controller = mock.Mock()

            q = DistUpgradeQuirks(mock_controller, mock.Mock())

            q._add_kms_overlay(boot_dir)
            self.assertFalse(os.path.exists(os.path.join(
                boot_dir, 'config.txt.distUpgrade')))

    def test_add_kms_overlay_no_changes(self):
        with tempfile.TemporaryDirectory() as boot_dir:
            boot_config = """\
arm_64bit=1
kernel=vmlinuz
initramfs initrd.img followkernel

# This line is implicitly in an [all] section and
# should prevent the quirk from doing anything
dtoverlay=vc4-kms-v3d,cma-128

[pi4]
max_framebuffers=2
"""
            with open(os.path.join(boot_dir, 'config.txt'), 'w') as f:
                f.write(boot_config)

            mock_controller = mock.Mock()
            q = DistUpgradeQuirks(mock_controller, mock.Mock())
            q._add_kms_overlay(boot_dir)

            self.assertFalse(os.path.exists(os.path.join(
                boot_dir, 'config.txt.distUpgrade')))
            with open(os.path.join(boot_dir, 'config.txt')) as f:
                self.assertTrue(f.read() == boot_config)

    def test_add_kms_overlay_with_changes(self):
        with tempfile.TemporaryDirectory() as boot_dir:
            config_txt = """\
arm_64bit=1
kernel=vmlinuz
initramfs initrd.img followkernel

[pi4]
max_framebuffers=2
"""
            expected_config_txt = """\
arm_64bit=1
kernel=vmlinuz
initramfs initrd.img followkernel

# added by do-release-upgrade (LP: #2065051)
dtoverlay=vc4-kms-v3d
disable_fw_kms_setup=1

[pi3+]
dtoverlay=vc4-kms-v3d,cma-128

[pi02]
dtoverlay=vc4-kms-v3d,cma-128

[all]
[pi4]
max_framebuffers=2
"""
            with open(os.path.join(boot_dir, 'config.txt'), 'w') as f:
                f.write(config_txt)

            mock_controller = mock.Mock()
            q = DistUpgradeQuirks(mock_controller, mock.Mock())
            q._add_kms_overlay(boot_dir)

            self.assertTrue(os.path.exists(os.path.join(
                boot_dir, 'config.txt.distUpgrade')))
            self.assertTrue(os.path.exists(os.path.join(
                boot_dir, 'config.txt')))
            with open(os.path.join(boot_dir, 'config.txt')) as f:
                self.assertTrue(f.read() == expected_config_txt)


class TestSnapQuirks(unittest.TestCase):

    def test_get_from_and_to_version(self):
        # Prepare the state for testing
        controller = mock.Mock()
        controller.fromDist = 'disco'
        controller.toDist = 'eoan'
        config = mock.Mock()
        q = DistUpgradeQuirks(controller, config)
        # Call method under test
        q._get_from_and_to_version()
        self.assertEqual(q._from_version, '19.04')
        self.assertEqual(q._to_version, '19.10')

    @mock.patch("subprocess.Popen", MockPopenSnap)
    def test_prepare_snap_replacement_data(self):
        # Prepare the state for testing
        controller = mock.Mock()
        config = mock.Mock()
        q = DistUpgradeQuirks(controller, config)
        q._from_version = "24.04"
        q._to_version = "24.10"
        # Call method under test

        controller.cache = {
            'ubuntu-desktop':
                make_mock_pkg(
                    name="ubuntu-desktop",
                    is_installed=True),
            'core18':
                make_mock_pkg(
                    name="core18",
                    is_installed=True),
            'gnome-3-28-1804':
                make_mock_pkg(
                    name="gnome-3-28-1804",
                    is_installed=True),
            'gtk-common-themes':
                make_mock_pkg(
                    name="gtk-common-themes",
                    is_installed=True),
            'gnome-calculator':
                make_mock_pkg(
                    name="gnome-calculator",
                    is_installed=True),
            'gnome-characters':
                make_mock_pkg(
                    name="gnome-characters",
                    is_installed=False),
            'gnome-logs':
                make_mock_pkg(
                    name="gnome-logs",
                    is_installed=False),
            'gnome-software':
                make_mock_pkg(
                    name="gnome-software",
                    is_installed=True),
            'snap-not-tracked':
                make_mock_pkg(
                    name="snap-not-tracked",
                    is_installed=True),
        }

        q._prepare_snap_replacement_data()
        # Check if the right snaps have been detected as installed and
        # needing refresh and which ones need installation
        self.maxDiff = None
        self.assertDictEqual(
            q._snap_list,
            {'core22': {
                'channel': 'stable',
                'command': 'install',
                'deb': None, 'snap-id': '1234'},
             'gnome-42-2204': {
                'channel': 'stable/ubuntu-24.10',
                'command': 'install',
                'deb': None, 'snap-id': '1234'},
             'snap-store': {
                'channel': '2/stable/ubuntu-24.10',
                'command': 'install',
                'deb': 'gnome-software',
                'snap-id': '1234'}})

    def test_is_deb2snap_metapkg_installed(self):
        # Prepare the state for testing
        controller = mock.Mock()
        config = mock.Mock()
        q = DistUpgradeQuirks(controller, config)
        q._from_version = "19.04"
        q._to_version = "19.10"
        controller.cache = {
            'ubuntu-desktop':
                make_mock_pkg(
                    name="ubuntu-desktop",
                    is_installed=True)
        }

        testdata = [
            # (input, expected output)
            ({}, False),
            ({'metapkg': None}, False),
            ({'metapkg': 'ubuntu-desktop'}, True),
            ({'metapkg': 'kubuntu-desktop'}, False),
            ({'metapkg': ['kubuntu-desktop', 'ubuntu-desktop']}, True),
            ({'metapkg': ['kubuntu-desktop', 'lubuntu-desktop']}, False),
        ]

        for data in testdata:
            self.assertEqual(q._is_deb2snap_metapkg_installed(data[0]),
                             data[1],
                             'Expected {1} for input {0}'.format(*data))

    @mock.patch("DistUpgrade.DistUpgradeQuirks.get_arch")
    @mock.patch("urllib.request.urlopen")
    def test_calculate_snap_size_requirements(self, urlopen, arch):
        # Prepare the state for testing
        arch.return_value = 'amd64'
        controller = mock.Mock()
        config = mock.Mock()
        q = DistUpgradeQuirks(controller, config)
        # We mock out _prepare_snap_replacement_data(), as this is tested
        # separately.
        q._prepare_snap_replacement_data = mock.Mock()
        q._snap_list = {
            'test-snap': {'command': 'install',
                          'deb': None, 'snap-id': '2',
                          'channel': 'stable/ubuntu-19.10'},
            'gnome-calculator': {'command': 'install',
                                 'deb': 'gnome-calculator',
                                 'snap-id': '1',
                                 'channel': 'stable/ubuntu-19.10'},
            'gnome-system-monitor': {'command': 'refresh',
                                     'channel': 'stable/ubuntu-19.10'}
        }
        q._to_version = "19.10"
        # Mock out urlopen in such a way that we get a mocked response based
        # on the parameters given but also allow us to check call arguments
        # etc.
        urlopen.side_effect = mock_urlopen_snap
        # Call method under test
        q._calculateSnapSizeRequirements()
        # Check if the size was calculated correctly
        self.assertEqual(q.extra_snap_space, 6218880)
        # Check if we only sent queries for the two command: install snaps
        self.assertEqual(urlopen.call_count, 2)
        # Make sure each call had the right headers and parameters
        for call in urlopen.call_args_list:
            req = call[0][0]
            self.assertIn(b"stable/ubuntu-19.10", req.data)
            self.assertDictEqual(
                req.headers,
                {'Snap-device-series': '16',
                 'Content-type': 'application/json',
                 'Snap-device-architecture': 'amd64'})

    @mock.patch("subprocess.run")
    def test_replace_debs_and_snaps(self, run_mock):
        controller = mock.Mock()
        config = mock.Mock()
        q = DistUpgradeQuirks(controller, config)
        q._snap_list = {
            'core18': {'command': 'refresh',
                       'channel': 'stable'},
            'gnome-3-28-1804': {'command': 'refresh',
                                'channel': 'stable/ubuntu-19.10'},
            'gtk-common-themes': {'command': 'refresh',
                                  'channel': 'stable/ubuntu-19.10'},
            'gnome-calculator': {'command': 'remove'},
            'gnome-characters': {'command': 'remove'},
            'gnome-logs': {'command': 'install',
                           'deb': 'gnome-logs',
                           'snap-id': '1234',
                           'channel': 'stable/ubuntu-19.10'},
            'gnome-system-monitor': {'command': 'install',
                                     'deb': 'gnome-system-monitor',
                                     'snap-id': '1234',
                                     'channel': 'stable/ubuntu-19.10'},
            'snap-store': {'command': 'install',
                           'deb': 'gnome-software',
                           'snap-id': '1234',
                           'channel': 'stable/ubuntu-19.10'}
        }
        q._to_version = "19.10"
        q._replaceDebsAndSnaps()
        # Make sure all snaps have been handled
        self.assertEqual(run_mock.call_count, 8)
        snaps_refreshed = {}
        snaps_installed = {}
        snaps_removed = []
        # Check if all the snaps that needed to be installed were installed
        # and those that needed a refresh - refreshed
        # At the same time, let's check that all the snaps were acted upon
        # while using the correct channel and branch
        for call in run_mock.call_args_list:
            args = call[0][0]
            if args[1] == 'refresh':
                snaps_refreshed[args[4]] = args[3]
            elif args[1] == 'install':
                snaps_installed[args[4]] = args[3]
            elif args[1] == 'remove':
                snaps_removed.append(args[2])
        self.assertDictEqual(
            snaps_refreshed,
            {'core18': 'stable',
             'gnome-3-28-1804': 'stable/ubuntu-19.10',
             'gtk-common-themes': 'stable/ubuntu-19.10'})
        self.assertDictEqual(
            snaps_installed,
            {'gnome-logs': 'stable/ubuntu-19.10',
             'gnome-system-monitor': 'stable/ubuntu-19.10',
             'snap-store': 'stable/ubuntu-19.10'})
        self.assertListEqual(
            snaps_removed,
            ['gnome-calculator',
             'gnome-characters'])
        # Make sure we marked the replaced ones for removal
        # Here we only check if the right number of 'packages' has been
        # added to the forced_obsoletes list - not all of those packages are
        # actual deb packages that will have to be removed during the upgrade
        self.assertEqual(controller.forced_obsoletes.append.call_count, 3)


if __name__ == "__main__":
    import logging
    logging.basicConfig(level=logging.DEBUG)
    unittest.main()
