#!/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 mock import patch
from DistUpgrade.DistUpgradeQuirks import DistUpgradeQuirks
from DistUpgrade.RiscvProfile import RiscvProfile

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)

    def test_disable_kdump_tools_on_install(self):
        """Test kdump-tools disabled if installed on upgrade."""
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        mock_pkg = make_mock_pkg(name="kdump-tools", is_installed=False)
        mock_pkg.marked_install = True
        mock_cache = {'kdump-tools': mock_pkg}

        with mock.patch(
            "DistUpgrade.DistUpgradeQuirks.subprocess.run"
        ) as popen_mock:
            q._disable_kdump_tools_on_install(mock_cache)

        expected_cmd = 'echo "kdump-tools kdump-tools/use_kdump boolean false"'
        expected_cmd += ' | debconf-set-selections'
        popen_mock.assert_called_once_with(expected_cmd, shell=True)

    def test_no_disable_kdump_tools_if_installed(self):
        """Test kdump-tools is not disabled on upgrade if already installed."""
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        mock_pkg = make_mock_pkg(name="kdump-tools", is_installed=True)
        mock_pkg.marked_install = True
        mock_cache = {'kdump-tools': mock_pkg}

        with mock.patch(
            "DistUpgrade.DistUpgradeQuirks.subprocess.run"
        ) as popen_mock:
            q._disable_kdump_tools_on_install(mock_cache)

        popen_mock.assert_not_called()

    def test_no_disable_kdump_tools_if_not_requested(self):
        """Test kdump-tools disable hook does nothing if not requested."""
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        # Do nothing if not found in cache
        mock_cache = {}

        with mock.patch(
            "DistUpgrade.DistUpgradeQuirks.subprocess.run"
        ) as popen_mock:
            q._disable_kdump_tools_on_install(mock_cache)

        popen_mock.assert_not_called()

        # Do nothing if not requested for install
        mock_pkg = make_mock_pkg(name="kdump-tools", is_installed=False)
        mock_pkg.marked_install = False
        mock_cache = {'kdump-tools': mock_pkg}

        with mock.patch(
            "DistUpgrade.DistUpgradeQuirks.subprocess.run"
        ) as popen_mock:
            q._disable_kdump_tools_on_install(mock_cache)

        popen_mock.assert_not_called()

    def test_test_and_fail_on_rva23(self):
        """
        Unit test for DistUpgradeQuirks.test_and_fail_on_rva23

        Expect abort for system not complying with RVA23U64.
        Expect no abort for compliant sytem.
        """

        # All RVA23U64 mandatory extensions
        good_cpu_info = RiscvProfile(
               'rv64imafdcbv_zic64b_zicbom_zicbop_zicboz_ziccamoa_ziccif_'
               'zicclsm_ziccrse_zicntr_zicond_zicsr_zihintntl_'
               'zihintpause_zihpm_zimop_za64rs_zawrs_zfa_zfhmin_zcb_'
               'zcmop_zvbb_zvfhmin_zvkt_zkt_supm')

        # Vector support missing
        bad_cpu_info = RiscvProfile(
               'rv64imafdcb_zic64b_zicbom_zicbop_zicboz_ziccamoa_ziccif_'
               'zicclsm_ziccrse_zicntr_zicond_zicsr_zihintntl_'
               'zihintpause_zihpm_zimop_za64rs_zawrs_zfa_zfhmin_zcb_'
               'zcmop_zkt_supm')

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

        with patch.object(RiscvProfile,
                          'read_cpuinfo') as mock_cpuinfo:
            mock_cpuinfo.return_value = good_cpu_info
            q._test_and_fail_on_rva23()
            self.assertEqual(0, len(mock_controller.abort.mock_calls))

        with patch.object(RiscvProfile,
                          'read_cpuinfo') as mock_cpuinfo:
            mock_cpuinfo.return_value = bad_cpu_info
            q._test_and_fail_on_rva23()
            self.assertEqual(1, len(mock_controller.abort.mock_calls))

    def test_test_and_fail_on_pi4_old_eeprom(self):
        from pathlib import Path
        with tempfile.TemporaryDirectory() as dt_dir:
            dt = Path(dt_dir)
            bootloader = dt / 'chosen' / 'bootloader'
            bootloader.mkdir(parents=True)

            mock_controller = mock.Mock()
            mock_controller.arch = 'arm64'

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

            # No d-t/compatible (no abort)
            mock_controller.abort.reset_mock()
            q._test_and_fail_on_pi_old_eeprom(dt_dir)
            self.assertFalse(mock_controller.abort.called)
            # d-t/compatible not rpi (no abort)
            mock_controller.abort.reset_mock()
            (dt / 'compatible').write_bytes(b'foo,bar\0foo,baz\0')
            q._test_and_fail_on_pi_old_eeprom(dt_dir)
            self.assertFalse(mock_controller.abort.called)
            # d-t/compatible is rpi 4, new build-timestamp (no abort)
            mock_controller.abort.reset_mock()
            (dt / 'compatible').write_bytes(
                b'raspberrypi,4-model-b\0brcm,bcm2711\0')
            (bootloader / 'build-timestamp').write_bytes(
                b'\x66\xf1\x67\x00')  # 1727096576 (2024-09-23)
            q._test_and_fail_on_pi_old_eeprom(dt_dir)
            self.assertFalse(mock_controller.abort.called)
            # d-t/compatible is rpi 4, old build-timestamp (abort)
            mock_controller.abort.reset_mock()
            (bootloader / 'build-timestamp').write_bytes(
                b'\x5f\x50\xdd\x7f')  # 1599135103 (2020-09-03)
            q._test_and_fail_on_pi_old_eeprom(dt_dir)
            self.assertTrue(mock_controller.abort.called)
            # d-t/compatible is rpi 5, old build-timestamp (abort)
            mock_controller.abort.reset_mock()
            (dt / 'compatible').write_bytes(
                b'raspberrypi,5-model-b\0brcm,bcm2712\0')
            (bootloader / 'build-timestamp').write_bytes(
                b'\x66\xf1\x67\x00')  # 1727096576 (2024-09-23)
            q._test_and_fail_on_pi_old_eeprom(dt_dir)
            self.assertTrue(mock_controller.abort.called)
            # d-t/compatible is rpi 5, no build-timestamp (abort)
            mock_controller.abort.reset_mock()
            (bootloader / 'build-timestamp').unlink()
            q._test_and_fail_on_pi_old_eeprom(dt_dir)
            self.assertTrue(mock_controller.abort.called)
            # d-t/compatible is rpi 3+, no build-timestamp (no abort)
            mock_controller.abort.reset_mock()
            (dt / 'compatible').write_bytes(
                b'raspberrypi,3-model-b-plus\0brcm,bcm2837\0')
            q._test_and_fail_on_pi_old_eeprom(dt_dir)
            self.assertFalse(mock_controller.abort.called)

    def test_and_fail_without_cgroupv2(self):
        mock_controller = mock.Mock()
        q = DistUpgradeQuirks(mock_controller, mock.Mock())

        m = mock.mock_open(read_data=(
            'sysfs /sys sysfs '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
            'proc /proc proc '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
            'udev /dev devtmpfs '
            'rw,nosuid,relatime,size=16029572k,mode=755,inode64 0 0\n'
            'devpts /dev/pts devpts '
            'rw,nosuid,noexec,relatime,gid=5,mode=620,ptmxmode=000 0 0\n'
            'tmpfs /run tmpfs '
            'rw,nosuid,nodev,noexec,relatime,mode=755,inode64 0 0\n'
            'securityfs /sys/kernel/security securityfs '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
            'tmpfs /dev/shm tmpfs '
            'rw,nosuid,nodev,inode64 0 0\n'
            'tmpfs /run/lock tmpfs '
            'rw,nosuid,nodev,noexec,relatime,size=5120k,inode64 0 0\n'
            'cgroup2 /sys/fs/cgroup cgroup2 '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
            'pstore /sys/fs/pstore pstore '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
        ))
        with mock.patch('DistUpgrade.DistUpgradeQuirks.open', m):
            q._test_and_fail_without_cgroupv2()
            self.assertEqual(0, len(mock_controller.abort.mock_calls))

        m = mock.mock_open(read_data=(
            'sysfs /sys sysfs '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
            'proc /proc proc '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
            'udev /dev devtmpfs '
            'rw,nosuid,relatime,size=16029572k,mode=755,inode64 0 0\n'
            'devpts /dev/pts devpts '
            'rw,nosuid,noexec,relatime,gid=5,mode=620,ptmxmode=000 0 0\n'
            'tmpfs /run tmpfs '
            'rw,nosuid,nodev,noexec,relatime,mode=755,inode64 0 0\n'
            'securityfs /sys/kernel/security securityfs '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
            'cgroup2 /sys/fs/cgroup/unified cgroup2 '
            'rw,nosuid,nodev,noexec,relatime,nsdelegate 0 0\n'
            'cgroup /sys/fs/cgroup/systemd cgroup '
            'rw,nosuid,nodev,noexec,relatime,xattr,name=systemd 0 0\n'
            'cgroup /sys/fs/cgroup/net_cls,net_prio cgroup '
            'rw,nosuid,nodev,noexec,relatime,net_cls,net_prio 0 0\n'
            'cgroup /sys/fs/cgroup/cpu,cpuacct cgroup '
            'rw,nosuid,nodev,noexec,relatime,cpu,cpuacct 0 0\n'
            'cgroup /sys/fs/cgroup/rdma cgroup '
            'rw,nosuid,nodev,noexec,relatime,rdma 0 0\n'
            'cgroup /sys/fs/cgroup/freezer cgroup '
            'rw,nosuid,nodev,noexec,relatime,freezer 0 0\n'
            'cgroup /sys/fs/cgroup/misc cgroup '
            'rw,nosuid,nodev,noexec,relatime,misc 0 0\n'
            'cgroup /sys/fs/cgroup/cpuset cgroup '
            'rw,nosuid,nodev,noexec,relatime,cpuset 0 0\n'
            'cgroup /sys/fs/cgroup/devices cgroup '
            'rw,nosuid,nodev,noexec,relatime,devices 0 0\n'
            'cgroup /sys/fs/cgroup/memory cgroup '
            'rw,nosuid,nodev,noexec,relatime,memory 0 0\n'
            'cgroup /sys/fs/cgroup/pids cgroup '
            'rw,nosuid,nodev,noexec,relatime,pids 0 0\n'
            'cgroup /sys/fs/cgroup/blkio cgroup '
            'rw,nosuid,nodev,noexec,relatime,blkio 0 0\n'
            'cgroup /sys/fs/cgroup/perf_event cgroup '
            'rw,nosuid,nodev,noexec,relatime,perf_event 0 0\n'
            'cgroup /sys/fs/cgroup/hugetlb cgroup '
            'rw,nosuid,nodev,noexec,relatime,hugetlb 0 0\n'
            'tmpfs /dev/shm tmpfs '
            'rw,nosuid,nodev,inode64 0 0\n'
            'tmpfs /run/lock tmpfs '
            'rw,nosuid,nodev,noexec,relatime,size=5120k,inode64 0 0\n'
            'pstore /sys/fs/pstore pstore '
            'rw,nosuid,nodev,noexec,relatime 0 0\n'
        ))
        with mock.patch('DistUpgrade.DistUpgradeQuirks.open', m):
            q._test_and_fail_without_cgroupv2()
            self.assertEqual(1, len(mock_controller.abort.mock_calls))

    def test_switch_to_systemd_coredump_kubuntu(self):
        """Test coredump switch happens for Kubuntu systems."""
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        kubuntu_pkg = make_mock_pkg(
            name="kubuntu-desktop", is_installed=True)
        kubuntu_pkg.marked_upgrade = False
        apport_pkg = make_mock_pkg(
            name="apport-core-dump-handler", is_installed=True)
        coredump_pkg = make_mock_pkg(
            name="systemd-coredump", is_installed=False)

        mock_cache = {
            'kubuntu-desktop': kubuntu_pkg,
            'apport-core-dump-handler': apport_pkg,
            'systemd-coredump': coredump_pkg,
        }
        q.controller.cache = mock_cache

        q._switch_to_systemd_coredump_for_kde_flavors()

        apport_pkg.mark_delete.assert_called_once_with(auto_fix=False)
        coredump_pkg.mark_install.assert_called_once_with(auto_fix=False)
        kubuntu_pkg.mark_upgrade.assert_called_once()

    def test_switch_to_dracut(self):
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        ubuntu_server_raspi = make_mock_pkg(
            name="ubuntu-server-raspi", is_installed=True)
        dracut = make_mock_pkg(
            name="dracut", is_installed=False)
        mock_cache = {
            'ubuntu-server-raspi': ubuntu_server_raspi,
            'dracut': dracut,
        }
        q.controller.cache = mock_cache

        q._force_dracut_install()
        dracut.mark_install.assert_called_once()

    def test_switch_to_piboot(self):
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        ubuntu_server_raspi = make_mock_pkg(
            name="ubuntu-server-raspi", is_installed=True)
        piboot_try = make_mock_pkg(
            name="piboot-try", is_installed=False)
        mock_cache = {
            'ubuntu-server-raspi': ubuntu_server_raspi,
            'piboot-try': piboot_try,
        }
        q.controller.cache = mock_cache

        q._force_piboot_try_install()
        piboot_try.mark_install.assert_called_once()

    def test_switch_to_systemd_coredump_ubuntustudio(self):
        """Test coredump switch happens for Ubuntu Studio systems."""
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        studio_pkg = make_mock_pkg(
            name="ubuntustudio-desktop-core", is_installed=True)
        studio_pkg.marked_upgrade = False
        apport_pkg = make_mock_pkg(
            name="apport-core-dump-handler", is_installed=True)
        coredump_pkg = make_mock_pkg(
            name="systemd-coredump", is_installed=False)

        mock_cache = {
            'ubuntustudio-desktop-core': studio_pkg,
            'apport-core-dump-handler': apport_pkg,
            'systemd-coredump': coredump_pkg,
        }
        q.controller.cache = mock_cache

        q._switch_to_systemd_coredump_for_kde_flavors()

        apport_pkg.mark_delete.assert_called_once_with(auto_fix=False)
        coredump_pkg.mark_install.assert_called_once_with(auto_fix=False)
        studio_pkg.mark_upgrade.assert_called_once()

    def test_no_switch_coredump_non_kde(self):
        """Test coredump switch is skipped for non-KDE flavors."""
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        apport_pkg = make_mock_pkg(
            name="apport-core-dump-handler", is_installed=True)
        coredump_pkg = make_mock_pkg(
            name="systemd-coredump", is_installed=False)

        mock_cache = {
            'apport-core-dump-handler': apport_pkg,
            'systemd-coredump': coredump_pkg,
        }
        q.controller.cache = mock_cache

        q._switch_to_systemd_coredump_for_kde_flavors()

        apport_pkg.mark_delete.assert_not_called()
        coredump_pkg.mark_install.assert_not_called()

    def test_switch_coredump_already_installed(self):
        """Test no-op when systemd-coredump is already installed."""
        q = DistUpgradeQuirks(mock.Mock(), mock.Mock())

        kubuntu_pkg = make_mock_pkg(
            name="kubuntu-desktop", is_installed=True)
        kubuntu_pkg.marked_upgrade = True
        coredump_pkg = make_mock_pkg(
            name="systemd-coredump", is_installed=True)

        mock_cache = {
            'kubuntu-desktop': kubuntu_pkg,
            'systemd-coredump': coredump_pkg,
        }
        q.controller.cache = mock_cache

        q._switch_to_systemd_coredump_for_kde_flavors()

        coredump_pkg.mark_install.assert_not_called()
        kubuntu_pkg.mark_upgrade.assert_not_called()


class TestSnapQuirks(unittest.TestCase):

    @mock.patch("subprocess.Popen", MockPopenSnap)
    def test_prepare_snap_replacement_data(self):
        # Prepare the state for testing
        controller = mock.Mock()
        config = mock.Mock()
        controller.fromVersion = '24.04'
        controller.toVersion = '24.10'
        q = DistUpgradeQuirks(controller, config)
        # 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,
            {'core24': {
                'channel': 'stable',
                'command': 'install',
                'deb': None, 'snap-id': '1234'},
             'gnome-46-2404': {
                'channel': 'stable/ubuntu-24.10',
                'command': 'install',
                'deb': None, 'snap-id': '1234'},
             'mesa-2404': {
                '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'},
             'thunderbird': {
                'channel': 'stable/ubuntu-24.10',
                'command': 'install',
                'deb': None,
                '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)
        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))

    def test_parse_deb2snap_json(self):
        controller = mock.Mock()
        config = mock.Mock()

        q = DistUpgradeQuirks(controller, config)
        q.controller.toVersion = '25.04'
        q._is_deb2snap_metapkg_installed = mock.Mock()
        q._is_deb2snap_metapkg_installed.return_value = True

        for from_version in ('22.04', '22.10'):
            q.controller.fromVersion = from_version

            m = mock.mock_open(read_data=(
                '{'
                '  "seeded": {'
                '    "test-snap-1": {'
                '      "metapkg": "ubuntu-desktop-minimal",'
                '      "from_channel": "latest/edge",'
                '      "to_channel": "latest/stable"'
                '    },'
                '    "test-snap-2": {'
                '      "deb": "test-deb-2",'
                '      "metapkg": "ubuntu-desktop-minimal"'
                '    },'
                '    "test-snap-3": {'
                '      "metapkg": "ubuntu-desktop-minimal",'
                '      "force_switch": true,'
                '      "from_channel": "1/stable/ubuntu-{FROM_VERSION}",'
                '      "to_channel": "2/stable/ubuntu-{TO_VERSION}"'
                '    }'
                '  },'
                '  "unseeded": {'
                '  }'
                '}'
            ))

            with mock.patch('DistUpgrade.DistUpgradeQuirks.open', m):
                seeded, unseeded = q._parse_deb2snap_json()

                self.assertEqual(unseeded, {})

                self.assertEqual(
                    seeded['test-snap-1'],
                    (
                        None,
                        'latest/edge',
                        'latest/stable',
                        False
                    )
                )
                self.assertEqual(
                    seeded['test-snap-2'],
                    (
                        'test-deb-2',
                        f'stable/ubuntu-{q.controller.fromVersion}',
                        f'stable/ubuntu-{q.controller.toVersion}',
                        False
                    )
                )
                self.assertEqual(
                    seeded['test-snap-3'],
                    (
                        None,
                        f'1/stable/ubuntu-{q.controller.fromVersion}',
                        f'2/stable/ubuntu-{q.controller.toVersion}',
                        True
                    )
                )

    @mock.patch("urllib.request.urlopen")
    def test_calculate_snap_size_requirements(self, urlopen):
        # Prepare the state for testing
        controller = mock.Mock()
        controller.arch = 'amd64'
        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'}
        }
        # 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._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()
