#!/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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure is not enabled.

                0 updates can be applied immediately.

                1 additional security update can be applied with ESM Infra.
                Learn more about enabling ESM Infra 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure 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 ESM Infra.
                Learn more about enabling ESM Infra 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure 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 ESM Infra to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo pro 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure is not enabled.

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

                Enable ESM Infra to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo pro 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure is not enabled.

                0 updates can be applied immediately.

                Enable ESM Infra to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo pro 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure is enabled.

                35 updates can be applied immediately.
                13 of these updates are ESM Infra 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure is enabled.

                47 updates can be applied immediately.
                13 of these updates are ESM Infra 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure 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, hide_esm_messages=False),
            textwrap.dedent(
                """
                Expanded Security Maintenance for Infrastructure is enabled.

                30 updates can be applied immediately.
                15 of these updates are ESM Infra security updates.
                15 of these updates are ESM Apps 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, hide_esm_messages=False),
            textwrap.dedent(
                """
                Expanded Security Maintenance for Infrastructure is not enabled.

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

                Enable ESM Infra to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo pro 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, hide_esm_messages=False),
            textwrap.dedent(
                """
                Expanded Security Maintenance for Infrastructure is not enabled.

                30 updates can be applied immediately.
                1 of these updates is an ESM Apps 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 ESM Infra.
                Learn more about enabling ESM Infra 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, hide_esm_messages=False),
            textwrap.dedent(
                """
                Expanded Security Maintenance for Applications is enabled.

                30 updates can be applied immediately.
                1 of these updates is an ESM Apps 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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Infrastructure 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 ESM Infra to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo pro 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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Applications 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 ESM Apps.
                Learn more about enabling ESM Apps 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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            textwrap.dedent(
                """\
                Expanded Security Maintenance for Applications is not enabled.

                0 updates can be applied immediately.

                Enable ESM Apps to receive additional future security updates.
                See https://ubuntu.com/esm or run: sudo pro 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, hide_esm_messages=False),
            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, hide_esm_messages=False),
            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_with_esm_sources_hide_esm(
        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, hide_esm_messages=True),
            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_for_lts_distro_not_in_esm_mode_yet_hide_esm(
        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, hide_esm_messages=True),
            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_esm_infra_disabled_wih_pkgs_and_esm_apps_on_lts_distro_hide_esm(
        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, hide_esm_messages=True),
            textwrap.dedent(
                """
                30 updates can be applied immediately.
                15 of these updates are standard security updates.
                To see these additional updates run: apt list --upgradable
                """).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_hide_esm(
        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, hide_esm_messages=True),
            textwrap.dedent(
                """
                30 updates can be applied immediately.
                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_and_esm_apps_enabled_hide_esm(
        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, hide_esm_messages=True),
            textwrap.dedent(
                """
                30 updates can be applied immediately.
                To see these additional updates run: apt list --upgradable
                """).lstrip())


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()
