Skip to content

Commit

Permalink
pkghelper: have alternative dry-run implementation
Browse files Browse the repository at this point in the history 
Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
  • Loading branch information
@ogayot
ogayot committed Sep 5, 2023
1 parent c09a32c commit 1670d07
Show file tree
Hide file tree
Showing 4 changed files with 122 additions and 84 deletions.
18 changes: 1 addition & 17 deletions subiquity/server/controllers/network.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,23 +108,7 @@ def wlan_support_install_state(self):
return self.app.package_installer.state_for_pkg("wpasupplicant")

async def _install_wpasupplicant(self):
if self.opts.dry_run:

async def fake_install(pkgname: str) -> PackageInstallState:
await asyncio.sleep(10 / self.app.scale_factor)
a = "DONE"
for k in self.app.debug_flags:
if k.startswith("wlan_install="):
a = k.split("=", 2)[1]
return getattr(PackageInstallState, a)

install_coro = fake_install
else:
install_coro = None

r = await self.app.package_installer.install_pkg(
"wpasupplicant", install_coro=install_coro
)
r = await self.app.package_installer.install_pkg("wpasupplicant")
log.debug("wlan_support_install_finished %s", r)
self._call_clients("wlan_support_install_finished", r)
if r == PackageInstallState.DONE:
Expand Down
75 changes: 50 additions & 25 deletions subiquity/server/pkghelper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
import asyncio
import logging
import os
from typing import Callable, Dict, Optional
from typing import Dict, List, Optional

import apt

Expand All @@ -26,9 +26,6 @@
log = logging.getLogger("subiquity.server.pkghelper")


InstallCoroutine = Optional[Callable[[str], PackageInstallState]]


class PackageInstaller:
"""Install packages from the pool on the ISO in the live session.
Expand All @@ -37,10 +34,9 @@ class PackageInstaller:
by the server installer.
"""

def __init__(self, app):
def __init__(self):
self.pkgs: Dict[str, asyncio.Task] = {}
self._cache: Optional[apt.Cache] = None
self.app = app

@property
def cache(self):
Expand All @@ -57,22 +53,12 @@ def state_for_pkg(self, pkgname: str) -> PackageInstallState:
else:
return PackageInstallState.INSTALLING

def start_installing_pkg(
self, pkgname: str, *, install_coro: InstallCoroutine = None
) -> None:
def start_installing_pkg(self, pkgname: str) -> None:
if pkgname not in self.pkgs:
if install_coro is None:
install_coro = self._install_pkg
self.pkgs[pkgname] = asyncio.create_task(install_coro(pkgname))

async def install_pkg(
self, pkgname: str, *, install_coro: InstallCoroutine = None
) -> PackageInstallState:
"""Install the requested package. The default implementation runs
apt-get (or will consider the package installed after two seconds in
dry-run mode). If a different implementation is wanted, one can specify
an alternative install coroutine"""
self.start_installing_pkg(pkgname, install_coro=install_coro)
self.pkgs[pkgname] = asyncio.create_task(self._install_pkg(pkgname))

async def install_pkg(self, pkgname: str) -> PackageInstallState:
self.start_installing_pkg(pkgname)
return await self.pkgs[pkgname]

async def _install_pkg(self, pkgname: str) -> PackageInstallState:
Expand All @@ -84,10 +70,6 @@ async def _install_pkg(self, pkgname: str) -> PackageInstallState:
if binpkg.installed:
log.debug("%s already installed", pkgname)
return PackageInstallState.DONE
if self.app.opts.dry_run:
log.debug("dry-run apt-get install %s", pkgname)
await asyncio.sleep(2 / self.app.scale_factor)
return PackageInstallState.DONE
if not binpkg.candidate.uri.startswith("cdrom:"):
log.debug(
"%s not available from cdrom (rather %s)", pkgname, binpkg.candidate.uri
Expand All @@ -107,3 +89,46 @@ async def _install_pkg(self, pkgname: str) -> PackageInstallState:
return PackageInstallState.DONE
else:
return PackageInstallState.FAILED


class DryRunPackageInstaller(PackageInstaller):
def __init__(self, app) -> None:
super().__init__()
self.scale_factor: float = app.scale_factor
self.debug_flags: List[str] = app.debug_flags
self.package_specific_impl = {
"wpasupplicant": self._install_wpa_supplicant,
}

async def _install_wpa_supplicant(self) -> PackageInstallState:
"""Special implementation for wpasupplicant (used by code related to
Wi-Fi interfaces)."""
await asyncio.sleep(10 / self.scale_factor)
status = "DONE"
for flag in self.debug_flags:
if flag.startswith("wlan_install="):
status = flag.split("=", 2)[1]
return getattr(PackageInstallState, status)

async def _install_pkg(self, pkgname: str) -> PackageInstallState:
if pkgname in self.package_specific_impl:
return await self.package_specific_impl[pkgname]()

log.debug("checking if %s is available", pkgname)
binpkg = self.cache.get(pkgname)
if not binpkg:
log.debug("%s not found", pkgname)
return PackageInstallState.NOT_AVAILABLE
if binpkg.installed:
log.debug("%s already installed", pkgname)
return PackageInstallState.DONE
log.debug("dry-run apt-get install %s", pkgname)
await asyncio.sleep(2 / self.scale_factor)
return PackageInstallState.DONE


def get_package_installer(app):
if app.opts.dry_run:
return DryRunPackageInstaller(app)
else:
return PackageInstaller()
4 changes: 2 additions & 2 deletions subiquity/server/server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@
from subiquity.server.dryrun import DRConfig
from subiquity.server.errors import ErrorController
from subiquity.server.geoip import DryRunGeoIPStrategy, GeoIP, HTTPGeoIPStrategy
from subiquity.server.pkghelper import PackageInstaller
from subiquity.server.pkghelper import get_package_installer
from subiquity.server.runner import get_command_runner
from subiquity.server.snapdapi import make_api_client
from subiquity.server.types import InstallerChannels
Expand Down Expand Up @@ -291,7 +291,7 @@ def __init__(self, opts, block_log_dir):
self.event_syslog_id = "subiquity_event.{}".format(os.getpid())
self.log_syslog_id = "subiquity_log.{}".format(os.getpid())
self.command_runner = get_command_runner(self)
self.package_installer = PackageInstaller(self)
self.package_installer = get_package_installer(self)

self.error_reporter = ErrorReporter(
self.context.child("ErrorReporter"), self.opts.dry_run, self.root
Expand Down
109 changes: 69 additions & 40 deletions subiquity/server/tests/test_pkghelper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,30 +18,36 @@

import attr

from subiquity.server.pkghelper import PackageInstaller, PackageInstallState
from subiquity.server.pkghelper import (
DryRunPackageInstaller,
PackageInstaller,
PackageInstallState,
)
from subiquity.server.pkghelper import log as PkgHelperLogger
from subiquitycore.tests import SubiTestCase
from subiquitycore.tests.mocks import make_app


class MockPackage:
@attr.s(auto_attribs=True)
class Candidate:
uri: str

def __init__(
self, installed: bool, name: str, candidate_uri: Optional[str] = None
) -> None:
self.installed = installed
self.name = name
if candidate_uri is None:
self.candidate = None
else:
self.candidate = self.Candidate(uri=candidate_uri)


@patch("apt.Cache", Mock(return_value={}))
class TestPackageInstaller(SubiTestCase):
class Package:
@attr.s(auto_attribs=True)
class Candidate:
uri: str

def __init__(
self, installed: bool, name: str, candidate_uri: Optional[str] = None
) -> None:
self.installed = installed
self.name = name
if candidate_uri is None:
self.candidate = None
else:
self.candidate = self.Candidate(uri=candidate_uri)

def setUp(self):
self.pkginstaller = PackageInstaller(make_app())
self.pkginstaller = PackageInstaller()

async def test_install_pkg_not_found(self):
self.assertEqual(
Expand All @@ -52,50 +58,73 @@ async def test_install_pkg_not_found(self):
async def test_install_pkg_already_installed(self):
with patch.dict(
self.pkginstaller.cache,
{"util-linux": self.Package(installed=True, name="util-linux")},
{"util-linux": MockPackage(installed=True, name="util-linux")},
):
self.assertEqual(
await self.pkginstaller.install_pkg("util-linux"),
PackageInstallState.DONE,
)

async def test_install_pkg_dr_install(self):
with patch.dict(
self.pkginstaller.cache,
{"python3-attr": self.Package(installed=False, name="python3-attr")},
):
with patch("subiquity.server.pkghelper.asyncio.sleep") as sleep:
self.assertEqual(
await self.pkginstaller.install_pkg("python3-attr"),
PackageInstallState.DONE,
)
sleep.assert_called_once()

async def test_install_pkg_not_from_cdrom(self):
with patch.dict(
self.pkginstaller.cache,
{
"python3-attr": self.Package(
"python3-attr": MockPackage(
installed=False,
name="python3-attr",
candidate_uri="http://archive.ubuntu.com",
)
},
):
with patch.object(self.pkginstaller.app.opts, "dry_run", False):
self.assertEqual(
await self.pkginstaller.install_pkg("python3-attr"),
PackageInstallState.NOT_AVAILABLE,
)


@patch("apt.Cache", Mock(return_value={}))
@patch("subiquity.server.pkghelper.asyncio.sleep")
class TestDryRunPackageInstaller(SubiTestCase):
def setUp(self):
app = make_app()
app.debug_flags = []
self.pkginstaller = DryRunPackageInstaller(app)

async def test_install_pkg(self, sleep):
with patch.dict(
self.pkginstaller.cache,
{"python3-attr": MockPackage(installed=False, name="python3-attr")},
):
with self.assertLogs(PkgHelperLogger, "DEBUG") as debug:
self.assertEqual(
await self.pkginstaller.install_pkg("python3-attr"),
PackageInstallState.NOT_AVAILABLE,
PackageInstallState.DONE,
)
sleep.assert_called_once()
self.assertIn(
"dry-run apt-get install %s", [record.msg for record in debug.records]
)

async def test_install_pkg_wpasupplicant_default_impl(self, sleep):
with patch.object(self.pkginstaller, "debug_flags", []):
self.assertEqual(
await self.pkginstaller.install_pkg("wpasupplicant"),
PackageInstallState.DONE,
)
sleep.assert_called_once()

async def test_install_pkg_alternative_impl(self):
async def impl(pkgname: str) -> PackageInstallState:
return PackageInstallState.FAILED
async def test_install_pkg_wpasupplicant_done(self, sleep):
with patch.object(self.pkginstaller, "debug_flags", ["wlan_install=DONE"]):
self.assertEqual(
await self.pkginstaller.install_pkg("wpasupplicant"),
PackageInstallState.DONE,
)
sleep.assert_called_once()

with patch.object(self.pkginstaller, "_install_pkg") as default_impl:
async def test_install_pkg_wpasupplicant_failed(self, sleep):
with patch.object(self.pkginstaller, "debug_flags", ["wlan_install=FAILED"]):
self.assertEqual(
await self.pkginstaller.install_pkg("python3-attr", install_coro=impl),
await self.pkginstaller.install_pkg("wpasupplicant"),
PackageInstallState.FAILED,
)

default_impl.assert_not_called()
sleep.assert_called_once()

0 comments on commit 1670d07

Please sign in to comment.