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

import apt_check
import io
import unittest
import textwrap

from unittest import mock


def get_message(*args, **kwds):
    with io.StringIO() as stream:
        apt_check.write_human_readable_summary(stream, *args, **kwds)
        return stream.getvalue()


class TestMotd(unittest.TestCase):
    """ Validate /etc/motd text """

    @mock.patch("apt_check.get_distro_version", return_value="16.04")
    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disabled_upto_date_esm_avail(
        self, _m_esm_distro, _m_is_lts, _m_distro_version
    ):
        self.assertEqual(
            get_message(upgrades=0, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=1,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                0 updates can be applied immediately.

                1 additional security update can be applied with UA Infra: ESM
                Learn more about enabling UA Infra: ESM service for Ubuntu 16.04 at
                https://ubuntu.com/16-04
                """))

    @mock.patch("apt_check.get_distro_version", return_value="18.04")
    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disabled_security_esm_avail(
        self, _m_esm_distro, _m_is_lts, _m_distro_version
    ):
        self.assertEqual(
            get_message(upgrades=15, security_updates=1,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=23,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                15 updates can be applied immediately.
                1 of these updates is a standard security update.
                To see these additional updates run: apt list --upgradable

                23 additional security updates can be applied with UA Infra: ESM
                Learn more about enabling UA Infra: ESM service for Ubuntu 18.04 at
                https://ubuntu.com/18-04
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disabled_security_no_esm_avail(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=15, security_updates=7,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                15 updates can be applied immediately.
                7 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable

                Enable UA Infra: ESM to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo ua status
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disabled_nosecurity(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=15, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                15 updates can be applied immediately.
                To see these additional updates run: apt list --upgradable

                Enable UA Infra: ESM to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo ua status
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disabled_noupdates(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=0, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                0 updates can be applied immediately.

                Enable UA Infra: ESM to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo ua status
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_enabled_nosecurity(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=35, security_updates=0,
                        esm_infra_updates=13, esm_apps_updates=0,
                        have_esm_infra=True, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is enabled.

                35 updates can be applied immediately.
                13 of these updates are UA Infra: ESM security updates.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_enabled_somesecurity(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=47, security_updates=7,
                        esm_infra_updates=13, esm_apps_updates=0,
                        have_esm_infra=True, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is enabled.

                47 updates can be applied immediately.
                13 of these updates are UA Infra: ESM security updates.
                7 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_enabled_noupdates(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=0, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=True, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=10),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is enabled.

                0 updates can be applied immediately.
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_and_esm_apps_enabled(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=0,
                        esm_infra_updates=15, esm_apps_updates=15,
                        have_esm_infra=True, have_esm_apps=True,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """
                UA Infra: Extended Security Maintenance (ESM) is enabled.

                30 updates can be applied immediately.
                15 of these updates are UA Infra: ESM security updates.
                15 of these updates are UA Apps: ESM security updates.
                To see these additional updates run: apt list --upgradable
                """).lstrip())

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disable_and_esm_apps_enabled(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=12,
                        esm_infra_updates=0, esm_apps_updates=15,
                        have_esm_infra=False, have_esm_apps=True,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                30 updates can be applied immediately.
                15 of these updates are UA Apps: ESM security updates.
                12 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable

                Enable UA Infra: ESM to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo ua status
                """).lstrip())

    @mock.patch("apt_check.get_distro_version", return_value="16.04")
    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disabled_wih_pkgs_and_esm_apps_enabled(
        self, _m_esm_distro, _m_is_lts, _m_distro_version
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=15,
                        esm_infra_updates=0, esm_apps_updates=1,
                        have_esm_infra=False, have_esm_apps=True,
                        disabled_esm_infra_updates=40,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                30 updates can be applied immediately.
                1 of these updates is a UA Apps: ESM security update.
                15 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable

                40 additional security updates can be applied with UA Infra: ESM
                Learn more about enabling UA Infra: ESM service for Ubuntu 16.04 at
                https://ubuntu.com/16-04
                """).lstrip())

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_esm_infra_disabled_wih_pkgs_and_esm_apps_enabled_on_lts_distro(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=15,
                        esm_infra_updates=0, esm_apps_updates=1,
                        have_esm_infra=False, have_esm_apps=True,
                        disabled_esm_infra_updates=40,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """
                UA Apps: Extended Security Maintenance (ESM) is enabled.

                30 updates can be applied immediately.
                1 of these updates is a UA Apps: ESM security update.
                15 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable
                """).lstrip())

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_esm_infra_disabled_and_esm_apps_disabled_with_pkgs(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=18,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=40),
            textwrap.dedent(
                """\
                UA Infra: Extended Security Maintenance (ESM) is not enabled.

                30 updates can be applied immediately.
                18 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable

                Enable UA Infra: ESM to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo ua status
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=True)
    def test_no_esm_infra_and_apps_index_in_esm_distro(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=18,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                30 updates can be applied immediately.
                18 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=False)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_no_esm_infra_and_apps_index_in_non_lts_distro(
        self, _m_esm_distro, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=18,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                30 updates can be applied immediately.
                18 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=False)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_esm_infra_disabled_and_esm_apps_disabled_with_no_esm_distro(
        self, _m_esm_distro, _m_is_lts,
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=18,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=40),
            textwrap.dedent(
                """\
                30 updates can be applied immediately.
                18 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=False)
    @mock.patch("subprocess.Popen")
    def test_message_for_distro_that_will_not_go_into_esm_mode(
        self, m_popen, _m_is_lts
    ):
        comm_mock = mock.MagicMock()
        comm_mock.communicate.return_value = (
            "(unknown)".encode("utf-8"), "")
        m_popen.return_value = comm_mock

        self.assertEqual(
            get_message(upgrades=30, security_updates=18,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=40),
            textwrap.dedent(
                """\
                30 updates can be applied immediately.
                18 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_for_lts_distro_not_in_esm_mode_yet(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=30, security_updates=18,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=10,
                        disabled_esm_apps_updates=40),
            textwrap.dedent(
                """\
                UA Apps: Extended Security Maintenance (ESM) is not enabled.

                30 updates can be applied immediately.
                18 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable

                40 additional security updates can be applied with UA Apps: ESM
                Learn more about enabling UA Apps: ESM service at https://ubuntu.com/esm
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=False)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_no_upgrades_non_lts_distro_without_esm_sources(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=0, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                0 updates can be applied immediately.
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=False)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_one_upgrade_non_lts_distro_without_esm_sources(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=1, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                1 update can be applied immediately.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=False)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_one_security_upgrade_non_lts_distro_without_esm_sources(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=1, security_updates=1,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                1 update can be applied immediately.
                1 of these updates is a standard security update.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_no_upgrades_lts_distro_without_esm_sources(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=0, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                0 updates can be applied immediately.
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_no_upgrades_lts_distro_with_esm_sources(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=0, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=False, have_esm_apps=False,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                UA Apps: Extended Security Maintenance (ESM) is not enabled.

                0 updates can be applied immediately.

                Enable UA Apps: ESM to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo ua status
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_one_upgrade_esm_distro_without_esm_sources(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=1, security_updates=0,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                1 update can be applied immediately.
                To see these additional updates run: apt list --upgradable
                """))

    @mock.patch("apt_check.is_lts_distro", return_value=True)
    @mock.patch("apt_check.is_esm_distro", return_value=False)
    def test_message_one_security_upgrade_esm_distro_without_esm_sources(
        self, _m_is_esm, _m_is_lts
    ):
        self.assertEqual(
            get_message(upgrades=1, security_updates=1,
                        esm_infra_updates=0, esm_apps_updates=0,
                        have_esm_infra=None, have_esm_apps=None,
                        disabled_esm_infra_updates=0,
                        disabled_esm_apps_updates=0),
            textwrap.dedent(
                """\
                1 update can be applied immediately.
                1 of these updates is a standard security update.
                To see these additional updates run: apt list --upgradable
                """))


class TestGetDistro(unittest.TestCase):
    @mock.patch("os.path.exists", return_value=True)
    def test_get_distro_happy_path(self, _m_path):
        read_data = (
            "UBUNTU_CODENAME=test\n"
            "NAME=Ubuntu\n"
            "VERSION=version\n"
            "ID=ubuntu\n"
            "ID_LIKE=debian\n"
            "PRETTY_NAME=Ubuntu XX.XX\n"
            "VERSION_ID=XX.XX\n"
        )

        expected_value = "test"
        with mock.patch(
            "builtins.open",
            mock.mock_open(read_data=read_data)
        ) as m_open:
            assert apt_check.get_distro() == expected_value
            assert 1 == _m_path.call_count
            assert [
                mock.call(apt_check.OS_RELEASE_PATH)
            ] == m_open.call_args_list

    @mock.patch("os.path.exists", return_value=True)
    def test_get_distro_no_ubuntu_codename(self, _m_path):
        read_data = (
            "NAME=Ubuntu\n"
            "VERSION=version\n"
            "ID=ubuntu\n"
            "ID_LIKE=debian\n"
            "PRETTY_NAME=Ubuntu XX.XX\n"
            "VERSION_ID=XX.XX\n"
        )
        expected_msg = (
            "Could not find UBUNTU_CODENAME in "
        ) + apt_check.OS_RELEASE_PATH

        with mock.patch(
            "builtins.open",
            mock.mock_open(read_data=read_data)
        ) as m_open:
            with self.assertRaises(Exception) as context:
                apt_check.get_distro()

            assert expected_msg == str(context.exception)

            assert 1 == _m_path.call_count
            assert [
                mock.call(apt_check.OS_RELEASE_PATH)
            ] == m_open.call_args_list

    @mock.patch("os.path.exists", return_value=False)
    def test_get_distro_no_os_release(self, _m_path):
        expected_msg = "File {} was not found on the system".format(
            apt_check.OS_RELEASE_PATH
        )

        with self.assertRaises(Exception) as context:
            apt_check.get_distro()

        assert expected_msg == str(context.exception)
        assert 1 == _m_path.call_count


class TestDistroVersion(unittest.TestCase):
    @mock.patch("os.path.exists", return_value=True)
    def test_get_distro_version_happy_path(self, _m_path):
        read_data = (
            "UBUNTU_CODENAME=test\n"
            "NAME=Ubuntu\n"
            "VERSION=version\n"
            "ID=ubuntu\n"
            "ID_LIKE=debian\n"
            "PRETTY_NAME=Ubuntu XX.XX\n"
            'VERSION_ID="XX.XX"\n'
        )

        expected_value = "XX.XX"
        with mock.patch(
            "builtins.open",
            mock.mock_open(read_data=read_data)
        ) as m_open:
            assert apt_check.get_distro_version() == expected_value
            assert 1 == _m_path.call_count
            assert [
                mock.call(apt_check.OS_RELEASE_PATH)
            ] == m_open.call_args_list


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