#!/usr/bin/python3

import os
import re
import shutil
import subprocess
import tempfile
import unittest

from aptsources.distro import get_distro
from aptsources.sourceslist import (SourceEntry, Deb822SourceEntry)
from copy import copy
from pathlib import Path
from types import SimpleNamespace
from urllib.request import urlopen
from urllib.error import (URLError, HTTPError)

from softwareproperties.cloudarchive import RELEASE_MAP
from softwareproperties.sourceutils import replace_source_entry


CODENAME = get_distro().codename

ADD_APT_REPOSITORY = shutil.which('add-apt-repository', path=f'.:{os.environ.get("PATH")}')

FAKE_URI = 'http://fake.archive.test.com/ubuntu'
SOURCESLIST_TEXT = f'deb {FAKE_URI} {CODENAME} main'

# types
TYPE_BINARY = 'deb'
TYPE_SOURCE = 'deb-src'
# pockets
RELEASE = 'release'
PROPOSED = 'proposed'
UPDATES = 'updates'
SECURITY = 'security'
BACKPORTS = 'backports'
# components
MAIN = frozenset(['main'])
UNIVERSE = frozenset(['universe'])
RESTRICTED = frozenset(['restricted'])
MULTIVERSE = frozenset(['multiverse'])

ARCHIVE_URI = 'http://fake.mirror.private.com/ubuntu/'
ARCHIVE_FILENAME_URI = re.sub(r'[:/.]+', '_', ARCHIVE_URI)
ARCHIVE_FILENAME = f'archive_uri-{ARCHIVE_FILENAME_URI}-{CODENAME}.list'

PPA_NAME = 'ubuntu-support-team/software-properties-autopkgtest'
PPA_URI = f'https://ppa.launchpadcontent.net/{PPA_NAME}/ubuntu/'
PPA_FILENAME = f'ubuntu-support-team-ubuntu-software-properties-autopkgtest-{CODENAME}.sources'
PPA_TRUSTED_FINGERPRINT = 'A17A D76F CBB7 A7D5 73C4  DF43 1E03 2FCE 2F88 6048'

ADD_APT_REPOSITORY_LONG_PARAMS = SimpleNamespace(
    remove='--remove',
    enable_source='--enable-source',
    component='--component',
    pocket='--pocket',
    yes='--yes',
    no_update='--no-update',
    ppa='--ppa',
    cloud='--cloud',
    uri='--uri',
    sourceslist='--sourceslist',
)

ADD_APT_REPOSITORY_SHORT_PARAMS = SimpleNamespace(
    remove='-r',
    enable_source='-s',
    component='-c',
    pocket='-p',
    yes='-y',
    no_update='-n',
    ppa='-P',
    cloud='-C',
    uri='-U',
    sourceslist='-S',
)


def has_network():
    try:
        with urlopen("https://launchpad.net/"):
            pass
    except (URLError, HTTPError):
        return False
    return True

def line(uri=None, binary=True, pocket=None, components=None):
    deb = TYPE_BINARY if binary else TYPE_SOURCE
    uri = uri or FAKE_URI
    pocket = f'-{pocket}' if pocket else ''
    components = ' '.join(components or MAIN)
    return f'{deb} {uri} {CODENAME}{pocket} {components}'

def deb822_lines(uri=None, binary=True, source=False, pocket=None, components=None):
    uris = uri or FAKE_URI

    suites = f'{CODENAME}' + (f'-{pocket}' if pocket else '')

    types = [TYPE_BINARY] if binary else []
    types += [TYPE_SOURCE] if source else []
    types = ' '.join(types)

    components = components or MAIN
    components = ' '.join(components)

    return [f'Types: {types}',f'URIs: {uris}', f'Suites: {suites}', f'Components: {components}']


@unittest.skipIf(ADD_APT_REPOSITORY is None, 'Could not find add-apt-repository script')
class TestAddAptRepository(unittest.TestCase):
    params = ADD_APT_REPOSITORY_LONG_PARAMS

    @classmethod
    def setUpClass(cls):
        # test apt.conf file
        d = tempfile.TemporaryDirectory()
        cls.addClassCleanup(d.cleanup)
        apt_conf = Path(d.name) / 'apt.conf'

        # set APT_CONFIG to the test apt.conf
        test_env = copy(os.environ)
        test_env['APT_CONFIG'] = str(apt_conf)

        # dir for test apt-get script
        d = tempfile.TemporaryDirectory()
        cls.addClassCleanup(d.cleanup)
        apt_get_dir = Path(d.name)

        # create a test apt-get script and prepend it to the PATH
        apt_get_output = apt_get_dir / 'test-apt-get-output'
        apt_get_script = apt_get_dir / 'apt-get'
        apt_get_script.touch(mode=0o755)
        apt_get_script.write_text(f'''#!/bin/bash\necho "$*" >> {apt_get_output}\n''')
        test_env['PATH'] = f'{apt_get_dir}:{test_env["PATH"]}'

        cls.apt_conf = apt_conf
        cls.apt_get_output = apt_get_output
        cls.test_env = test_env

    def setUp(self):
        d = tempfile.TemporaryDirectory()
        self.addCleanup(d.cleanup)
        dir_etc = Path(d.name)
        self.apt_conf.write_text(f'Dir::Etc "{dir_etc}";')

        dir_etc.joinpath('apt.conf.d').mkdir(exist_ok=True)

        self.sourceslist = dir_etc / 'sources.list'
        self.sourceslist.touch()
        self.sourceslist.write_text(SOURCESLIST_TEXT)

        self.sourceslistd = dir_etc / 'sources.list.d'
        self.sourceslistd.mkdir(exist_ok=True)

        self.trustedgpgd = dir_etc / 'trusted.gpg.d'
        self.trustedgpgd.mkdir(exist_ok=True)

        self._sourcesfile = None
        self._trustedfile = None

        self.test_env.pop('LC_ALL', None)

        self.testing_ppa = False

        self.testing_uca = False
        self.test_env.pop('CA_ALLOW_CODENAME', None)

        self.noupdate = False
        self.pocket = None
        self.components = None
        self.enable_source = 0

    @property
    def sourcesfile(self):
        return self._sourcesfile

    @property
    def trustedfile(self):
        return self._trustedfile

    @sourcesfile.setter
    def sourcesfile(self, filename):
        self._sourcesfile = self._file_setter(self.sourceslistd, filename)

    @trustedfile.setter
    def trustedfile(self, filename):
        self._trustedfile = self._file_setter(self.trustedgpgd, filename)

    def _file_setter(self, d, f):
        if not f:
            return None
        path = d / f
        self.addCleanup(path.unlink, missing_ok=True)
        return path

    @property
    def parameter_enable_source(self):
        return [self.params.enable_source] * self.enable_source

    @property
    def parameter_pocket(self):
        return [self.params.pocket, self.pocket] if self.pocket else []

    @property
    def parameter_components(self):
        params = []
        for c in self.components or []:
            params.append(self.params.component)
            params.append(c)
        return params

    @property
    def parameter_all(self):
        return self.parameter_enable_source + self.parameter_pocket + self.parameter_components

    def line(self, uri=None, binary=True, pocket=None, components=None):
        return line(uri=uri or ARCHIVE_URI,
                    binary=binary,
                    pocket=pocket or self.pocket,
                    components=components or self.components)

    @unittest.skipUnless(os.geteuid() == 0, 'requires root')
    def call_add_apt_repository(self, args):
        self.apt_get_output.unlink(missing_ok=True)

        n = [self.params.no_update] if self.noupdate else []
        result = subprocess.run([ADD_APT_REPOSITORY] + ['-y'] + n + args,
                                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                env=self.test_env)

        if result.returncode != 0:
            self.fail(f'add-apt-repository failed: {result.stdout}')

    def check_sourceentry(self, expected_line, content_lines):
        expected_entry = SourceEntry(expected_line)
        expected_entry.comps = set(expected_entry.comps)

        for (index, line) in enumerate(content_lines):
            entry = SourceEntry(line)
            entry.comps = set(entry.comps)
            if entry == expected_entry:
                return content_lines[:index] + content_lines[index+1:]

        content = '\n'.join(content_lines)
        self.fail(f"Missing line: '{expected_line}' in: '{content}'")

    def check_expected_lines(self, expected_lines, lines):
        for expected_line in expected_lines:
            lines = self.check_sourceentry(expected_line, lines)

        # ignore invalid or disabled lines
        lines = list(filter(lambda l: not SourceEntry(l).invalid, lines))
        lines = list(filter(lambda l: not SourceEntry(l).disabled, lines))
        if lines:
            unexpected_lines = '\n'.join(lines)
            self.fail(f"Unexpected lines:\n{unexpected_lines}")

    def check_expected_deb822_lines(self, expected_lines, lines):
        def filter_lines(lines):
            in_ignored_field = False
            for line in lines:
                if line.startswith('Signed-By:'):
                    in_ignored_field = True
                    continue

                if in_ignored_field and line.startswith(' '):
                    continue

                in_ignored_field = False
                yield line

        # Ignore lines that should be there, but we are not checking for, e.g.
        # Signed-By.
        lines = filter_lines(lines)
        missing_lines = list(set(expected_lines) - set(lines))
        unexpected_lines = list(set(lines) - set(expected_lines))


        if missing_lines:
            self.fail(f'Missing lines: {missing_lines} in: {lines}')

        if unexpected_lines:
            self.fail(f'Unexpected lines: {unexpected_lines} in: {lines}')

    def apt_get_output_lines(self):
        try:
            return self.apt_get_output.read_text().splitlines()
        except FileNotFoundError:
            return []

    def check_apt_get_update(self):
        if self.noupdate:
            self.assertNotIn('update', self.apt_get_output_lines())
        else:
            self.assertIn('update', self.apt_get_output_lines())

    def check_uca_keyring(self):
        if self.testing_uca:
            self.assertIn('install -y ubuntu-cloud-keyring', self.apt_get_output_lines())

    def check_ppa_gpgkey(self):
        if not self.testing_ppa:
            return

        with open(self.sourcesfile, mode='r', encoding='utf-8') as f:
            entry = Deb822SourceEntry(f.read(), '')

        keydata = entry.section.get('Signed-By')
        keydata = keydata.splitlines()
        keydata = [l.lstrip() for l in keydata]
        keydata = [l if l != '.' else '' for l in keydata]
        keydata = '\n'.join(keydata)

        with tempfile.NamedTemporaryFile(mode='w') as f:
            f.write(keydata)
            f.flush()

            with tempfile.TemporaryDirectory() as homedir:
                Path(homedir).chmod(0o700) # so gpg doesn't complain to stderr
                cmd = f'gpg -q --homedir {homedir} --import {f.name}'
                subprocess.check_call(cmd.split())
                keyring = os.path.join(homedir, 'pubring.kbx')
                cmd = f'gpg -q --homedir {homedir} --no-default-keyring --keyring {keyring} --fingerprint'
                result = subprocess.run(cmd.split(), stdout=subprocess.PIPE, encoding='UTF-8')
                self.assertIn(PPA_TRUSTED_FINGERPRINT, result.stdout)

    def check_file(self, testfile, exists=True):
        if not testfile:
            return
        msg = f'{"missing" if exists else "unexpected"} {testfile}'
        self.assertEqual(testfile.exists(), exists, msg)

    def run_sourceslistd_add_remove_test(self, args, expected_lines=[]):
        self.check_file(self.sourcesfile, exists=False)
        self.check_file(self.trustedfile, exists=False)

        self.call_add_apt_repository(args)

        self.check_file(self.sourcesfile)
        self.check_file(self.trustedfile)

        if self.testing_ppa:
            self.check_expected_deb822_lines(expected_lines, self.sourcesfile.read_text().splitlines())
        else:
            self.check_expected_lines(expected_lines, self.sourcesfile.read_text().splitlines())

        self.check_apt_get_update()
        self.check_ppa_gpgkey()
        self.check_uca_keyring()

        self.call_add_apt_repository([self.params.remove] + args)

        self.check_file(self.sourcesfile, exists=False)

        # trusted.gpg.d file isn't removed by add-apt-repository
        self.check_file(self.trustedfile)

    def run_sourceslistd_test(self, args, expected_line):
        if self.testing_ppa:
            expected_lines = deb822_lines(uri=PPA_URI)
        else:
            source_entry = replace_source_entry(SourceEntry(expected_line), type=TYPE_SOURCE)
            source_entry.set_enabled(self.enable_source > 0)

            expected_lines = [expected_line, str(source_entry)]
        test_args = self.parameter_all
        test_args += args
        self.run_sourceslistd_add_remove_test(test_args, expected_lines)

    def test_uri(self):
        self.sourcesfile = ARCHIVE_FILENAME
        args = [self.params.uri, ARCHIVE_URI]
        self.run_sourceslistd_test(args, self.line())

    def test_uri_noupdate(self):
        self.noupdate = True
        self.test_uri()

    def test_sourceslist(self):
        self.sourcesfile = ARCHIVE_FILENAME
        line = self.line()
        args = [self.params.sourceslist, line]
        self.run_sourceslistd_test(args, line)

    def test_sourceslist_noupdate(self):
        self.noupdate = True
        self.test_sourceslist()

    @unittest.skipUnless(has_network(), 'requires network')
    def test_ppa(self, /, with_param=True, prefix=False):
        self.testing_ppa = True

        self.sourcesfile = PPA_FILENAME
        self.trustedfile = None
        args = []
        if with_param:
            args += [self.params.ppa]
        if prefix:
            args += [f'ppa:{PPA_NAME}']
        else:
            args += [PPA_NAME]
        self.run_sourceslistd_test(args, None)

    def test_ppa_noupdate(self):
        self.noupdate = True
        self.test_ppa()

    def test_ppa_with_prefix(self):
        self.test_ppa(prefix=True)

    def test_ppa_without_param_with_prefix(self):
        self.test_ppa(with_param=False, prefix=True)

    def _test_uca(self, uca_release, /, with_param, prefix):
        self.testing_uca = True
        if CODENAME != RELEASE_MAP.get(uca_release):
            self.test_env['CA_ALLOW_CODENAME'] = CODENAME

        self.sourcesfile = f'cloudarchive-{uca_release}.list'
        line = f'deb http://ubuntu-cloud.archive.canonical.com/ubuntu {CODENAME}-updates/{uca_release} main'
        args = []
        if with_param:
            args += [self.params.cloud]
        if prefix:
            args += [f'{prefix}:{uca_release}']
        else:
            args += [uca_release]
        self.run_sourceslistd_test(args, line)

    def test_uca(self, /, with_param=True, prefix=None):
        uca_releases = list(RELEASE_MAP.keys())
        self.assertGreater(len(uca_releases), 0)
        for uca_release in uca_releases:
            with self.subTest(uca_release=uca_release):
                self.setUp()
                self._test_uca(uca_release, with_param=with_param, prefix=prefix)
                self.tearDown()

    def test_uca_noupdate(self):
        self.noupdate = True
        self.test_uca()

    def test_uca_with_prefix_uca(self):
        self.test_uca(prefix='uca')

    def test_uca_with_prefix_uca_without_param(self):
        self.test_uca(with_param=False, prefix='uca')

    def test_uca_with_prefix_cloud_archive(self):
        self.test_uca(prefix='cloud-archive')

    def test_uca_with_prefix_cloud_archive_without_param(self):
        self.test_uca(with_param=False, prefix='cloud-archive')

    class TestLinesGroup(object):
        def __init__(self, /, test_pocket=None, test_components=frozenset(), repo_n=1):
            self.test_pocket = test_pocket
            self.test_components = test_components
            self.repo_n = repo_n

            self.lines = []
            self.expected_lines = []
            self.after_remove = []

        def next_repo(self):
            self.repo_n += 1

        def entry(self, /, binary=True, pocket=None, components=MAIN):
            return SourceEntry(self.line(binary=binary, pocket=pocket, components=components))

        def line(self, /, binary=True, pocket=None, components=MAIN):
            n = self.repo_n
            return line(uri=f'{FAKE_URI}/source{n}', binary=binary, pocket=pocket, components=components)

    def check_lines_group(self, args, group):
        self.sourceslist.write_text('\n'.join(group.lines))

        self.call_add_apt_repository(args)
        self.check_expected_lines(group.expected_lines,
                                  self.sourceslist.read_text().splitlines())


        self.call_add_apt_repository([self.params.remove] + args)
        self.check_expected_lines(group.after_remove,
                                  self.sourceslist.read_text().splitlines())

    def add_lines_comp(self, group, components, first_line=True, other_lines=frozenset()):
        group.lines.append(group.line(components=components))

        expected = copy(components)
        if first_line:
            # all new components will be added to (only) the first line
            expected |= group.test_components
            expected -= other_lines
        group.expected_lines.append(group.line(components=expected))

        after_remove = components - group.test_components
        if after_remove:
            group.after_remove.append(group.line(components=after_remove))

    def test_component(self):
        # with no self.components, we test the no-prefix case using MAIN
        test_components = copy(self.components or MAIN)
        args = self.parameter_components or list(test_components)

        group = self.TestLinesGroup(test_components=test_components)

        # source1, with two lines, one component on each
        self.add_lines_comp(group, MAIN, other_lines=UNIVERSE)
        self.add_lines_comp(group, UNIVERSE, first_line=False)

        # source2, one line with two components
        group.next_repo()
        self.add_lines_comp(group, MAIN | UNIVERSE)

        # source3, same as source2 but different components
        group.next_repo()
        self.add_lines_comp(group, RESTRICTED | MULTIVERSE)

        # source4, one line, one component
        group.next_repo()
        self.add_lines_comp(group, RESTRICTED)

        self.check_lines_group(args, group)

    def test_component_main(self):
        self.components = MAIN
        self.test_component()

    def test_component_main_noupdate(self):
        self.noupdate = True
        self.test_component_main()

    def test_component_universe(self):
        self.components = UNIVERSE
        self.test_component()

    def test_component_universe_noupdate(self):
        self.noupdate = True
        self.test_component_universe()

    def test_component_restricted(self):
        self.components = RESTRICTED
        self.test_component()

    def test_component_multiverse(self):
        self.components = MULTIVERSE
        self.test_component()

    def test_component_main_universe(self):
        self.components = MAIN | UNIVERSE
        self.test_component()

    def test_component_universe_multiverse(self):
        self.components = UNIVERSE | MULTIVERSE
        self.test_component()

    def test_component_all(self):
        self.components = MAIN | UNIVERSE | RESTRICTED | MULTIVERSE
        self.test_component()

    def test_component_all_except_main(self):
        self.components = UNIVERSE | RESTRICTED | MULTIVERSE
        self.test_component()

    def add_lines_pocket(self, group, pocket, components=MAIN, other_pockets=[], other_comps=frozenset(), first_line=True):
        line = group.line(pocket=pocket, components=components)
        group.lines.append(line)
        group.expected_lines.append(line)
        if (pocket or RELEASE) != group.test_pocket:
            group.after_remove.append(line)

    def _test_pocket(self):
        self.assertIsNotNone(self.pocket)

        args = self.parameter_pocket

        group = self.TestLinesGroup(test_pocket=self.pocket)

        # source1, with two lines, one pocket, one component on each
        self.add_lines_pocket(group, None, MAIN)
        self.add_lines_pocket(group, None, UNIVERSE)
        if self.pocket not in [RELEASE]:
            group.expected_lines.append(group.line(pocket=self.pocket, components=MAIN | UNIVERSE))

        # source2, with two lines, two pockets, varying components on each
        group.next_repo()
        self.add_lines_pocket(group, None, MAIN | UNIVERSE)
        self.add_lines_pocket(group, UPDATES, MAIN)
        if self.pocket not in [RELEASE, UPDATES]:
            group.expected_lines.append(group.line(pocket=self.pocket, components=MAIN | UNIVERSE))

        # source3, with three lines, two pockets, varying components on each
        group.next_repo()
        self.add_lines_pocket(group, None, MAIN)
        self.add_lines_pocket(group, None, UNIVERSE)
        self.add_lines_pocket(group, UPDATES, MAIN)
        if self.pocket not in [RELEASE, UPDATES]:
            group.expected_lines.append(group.line(pocket=self.pocket, components=MAIN | UNIVERSE))

        # source4, no release pocket
        group.next_repo()
        self.add_lines_pocket(group, UPDATES, MAIN)
        if self.pocket not in [UPDATES]:
            group.expected_lines.append(group.line(pocket=self.pocket, components=MAIN | RESTRICTED))

        # source5, all pockets exist
        group.next_repo()
        self.add_lines_pocket(group, None, MAIN)
        self.add_lines_pocket(group, UPDATES, MAIN)
        self.add_lines_pocket(group, PROPOSED, MAIN)
        self.add_lines_pocket(group, SECURITY, MAIN)

        # source6, release lines specified with and without pocket
        group.next_repo()
        self.add_lines_pocket(group, None, MAIN)
        self.add_lines_pocket(group, RELEASE, MULTIVERSE)
        self.add_lines_pocket(group, UPDATES, MAIN)
        if self.pocket not in [RELEASE, UPDATES]:
            group.expected_lines.append(group.line(pocket=self.pocket, components=MAIN | MULTIVERSE))

        self.check_lines_group(args, group)

    def test_pocket_release(self):
        self.pocket = RELEASE
        self._test_pocket()

    def test_pocket_release_noupdate(self):
        self.noupdate = True
        self.test_pocket_release()

    def test_pocket_updates(self):
        self.pocket = UPDATES
        self._test_pocket()

    def test_pocket_updates_noupdate(self):
        self.noupdate = True
        self.test_pocket_updates()

    def test_pocket_proposed(self):
        self.pocket = PROPOSED
        self._test_pocket()

    def test_pocket_security(self):
        self.pocket = SECURITY
        self._test_pocket()

    def add_lines_es(self, group, pocket, components, /, binary=True, disabled=False, alone=False, plus_comps=[], minus_comps=False):
        binary_entry = replace_source_entry(group.entry(pocket=pocket, components=components), disabled=disabled)
        enabled_source_entry = group.entry(binary=False, pocket=pocket, components=components)
        disabled_source_entry = replace_source_entry(enabled_source_entry, disabled=True)

        if binary:
            binary_line = str(binary_entry)
            group.lines.append(binary_line)
            group.expected_lines.append(binary_line)
            group.after_remove.append(binary_line)
            # with binary, 'alone' indicates no matching source line exists,
            # so we need to create one if enable_source > 1
            # if minus_comps is True, our comps will be added to another line
            if alone and not disabled and not minus_comps and self.enable_source > 1:
                group.expected_lines.append(str(enabled_source_entry))
                group.after_remove.append(str(disabled_source_entry))
        else:
            enabled_line = str(enabled_source_entry)
            disabled_line = str(disabled_source_entry)
            if not disabled:
                group.lines.append(enabled_line)
            else:
                group.lines.append(disabled_line)
            if plus_comps and self.enable_source > 1:
                # include comps from other binary-only lines
                comps = set(enabled_source_entry.comps) | set(plus_comps)
                enabled_line = str(replace_source_entry(enabled_source_entry, comps=comps))
                disabled_line = str(replace_source_entry(disabled_source_entry, comps=comps))
            if disabled and alone:
                # with source, 'alone' indicates no matching (enabled) binary line exists,
                # so the source line will *not* be enabled
                group.expected_lines.append(disabled_line)
            else:
                # if already enabled, or if not 'alone', we expect it to be enabled
                group.expected_lines.append(enabled_line)
            group.after_remove.append(disabled_line)

    def add_lines_es_binary(self, group, pocket, components, /, **kwargs):
        self.add_lines_es(group, pocket, components, binary=True, **kwargs)

    def add_lines_es_source(self, group, pocket, components, /, **kwargs):
        self.add_lines_es(group, pocket, components, binary=False, **kwargs)

    def add_lines_es_both(self, group, pocket, components, /, source_disabled=False, **kwargs):
        self.add_lines_es_binary(group, pocket, components, **kwargs)
        if kwargs.get('disabled'):
            # if both lines are disabled, source line is effectively 'alone'
            # as there is no enabled matching binary line
            kwargs['alone'] = True
        if source_disabled:
            kwargs['disabled'] = True
        self.add_lines_es_source(group, pocket, components, **kwargs)

    def _test_enable_source(self):
        self.assertGreater(self.enable_source, 0)

        args = self.parameter_enable_source

        group = self.TestLinesGroup()

        # source1, simple test with just binary lines
        self.add_lines_es_binary(group, RELEASE, MAIN | UNIVERSE, alone=True)
        self.add_lines_es_binary(group, UPDATES, MAIN | UNIVERSE, alone=True)
        self.add_lines_es_binary(group, PROPOSED, MAIN | UNIVERSE, alone=True)
        self.add_lines_es_binary(group, SECURITY, MAIN | UNIVERSE, alone=True)

        # source2, all with matching disabled source lines
        group.next_repo()
        self.add_lines_es_both(group, None, MAIN | UNIVERSE, source_disabled=True)
        self.add_lines_es_both(group, None, RESTRICTED | MULTIVERSE, source_disabled=True)
        self.add_lines_es_both(group, UPDATES, MAIN | UNIVERSE | RESTRICTED | MULTIVERSE, source_disabled=True)
        self.add_lines_es_both(group, SECURITY, MAIN | UNIVERSE, source_disabled=True)

        # source3, mix of just binary lines plus some disabled and enabled source lines
        group.next_repo()
        self.add_lines_es_both(group, None, MAIN | UNIVERSE, source_disabled=True, plus_comps=RESTRICTED | MULTIVERSE)
        self.add_lines_es_binary(group, None, RESTRICTED | MULTIVERSE, alone=True, minus_comps=True)
        self.add_lines_es_both(group, UPDATES, MAIN | UNIVERSE | RESTRICTED | MULTIVERSE, source_disabled=True)
        self.add_lines_es_both(group, SECURITY, MAIN | UNIVERSE)

        # source4, disabled binary lines to verify disabled source lines aren't enabled
        group.next_repo()
        self.add_lines_es_both(group, None, MAIN | UNIVERSE, disabled=True)
        self.add_lines_es_both(group, UPDATES, MULTIVERSE, disabled=True)

        self.check_lines_group(args, group)

    def test_enable_source_1(self):
        self.enable_source = 1
        self._test_enable_source()

    def test_enable_source_1_noupdate(self):
        self.no_update = True
        self.test_enable_source_1()

    def test_enable_source_2(self):
        self.enable_source = 2
        self._test_enable_source()

    def test_enable_source_2_noupdate(self):
        self.no_update = True
        self.test_enable_source_2()

    @classmethod
    def add_locale_test(cls, testname, locale):
        newname = f"{testname}_locale_{re.sub(r'[-.]+', '_', locale)}"
        if hasattr(cls, newname):
            return
        def _test_locale(self):
            self.test_env['LC_ALL'] = locale
            getattr(self, testname)()
        setattr(cls, newname, _test_locale)

    @classmethod
    def add_combination_test(cls, testname, /, enable_source=0, pocket=None, components=None):
        newname = testname
        if enable_source > 0:
            newname += f'_enable_source_{enable_source}'
        if pocket:
            newname += f'_pocket_{pocket}'
        if components:
            newname += f'_components_{"_".join(components)}'
        if hasattr(cls, newname):
            return
        def _test_combination(self):
            self.enable_source = enable_source
            self.pocket = pocket
            self.components = components
            getattr(self, testname)()
        setattr(cls, newname, _test_combination)


class TestAddAptRepositoryShortParams(TestAddAptRepository):
    params = ADD_APT_REPOSITORY_SHORT_PARAMS


# Add various combination tests
for test in ['uri', 'sourceslist']:
    for locale in ['C', 'C.UTF-8']:
        TestAddAptRepository.add_locale_test(f'test_{test}', locale=locale)
    for enable_source in [0, 1, 2]:
        for pocket in [None, RELEASE, UPDATES, PROPOSED, SECURITY]:
            for components in [None, MAIN, UNIVERSE, RESTRICTED, MULTIVERSE, MAIN | UNIVERSE,
                               MAIN | UNIVERSE | RESTRICTED | MULTIVERSE]:
                TestAddAptRepository.add_combination_test(f'test_{test}',
                                                          enable_source=enable_source,
                                                          pocket=pocket,
                                                          components=components)


if __name__ == "__main__":
    unittest.main(verbosity=2)
