From b8eea6ae7a3e18d0c06686fb08fecdcb1784bb20 Mon Sep 17 00:00:00 2001
From: Alberto Contreras <alberto.contreras@canonical.com>
Date: Tue, 24 Sep 2024 17:35:07 +0200
Subject: [PATCH] feat(integration_test): add CLOUD_INIT_PKG setting

Add CLOUD_INIT_PKG to control what specific cloud-init package to
install when CLOUD_INIT_SOURCE is PROPOSED, UPGRADE or points to a ppa.

SC-1858
---
 tests/integration_tests/instances.py          | 50 +++++++++++++------
 .../integration_tests/integration_settings.py |  4 ++
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/tests/integration_tests/instances.py b/tests/integration_tests/instances.py
index 1c8344ab916..a5ce8f1e14f 100644
--- a/tests/integration_tests/instances.py
+++ b/tests/integration_tests/instances.py
@@ -181,15 +181,16 @@ def install_new_cloud_init(
         self,
         source: CloudInitSource,
         clean=True,
+        pkg: str = integration_settings.CLOUD_INIT_PKG,
     ):
         if source == CloudInitSource.DEB_PACKAGE:
             self.install_deb()
         elif source == CloudInitSource.PPA:
-            self.install_ppa()
+            self.install_ppa(pkg)
         elif source == CloudInitSource.PROPOSED:
-            self.install_proposed_image()
+            self.install_proposed_image(pkg)
         elif source == CloudInitSource.UPGRADE:
-            self.upgrade_cloud_init()
+            self.upgrade_cloud_init(pkg)
         else:
             raise RuntimeError(
                 "Specified to install {} which isn't supported here".format(
@@ -201,8 +202,8 @@ def install_new_cloud_init(
         if clean:
             self.instance.clean()
 
-    def install_proposed_image(self):
-        log.info("Installing proposed image")
+    def install_proposed_image(self, pkg: str):
+        log.info("Installing %s from -proposed", pkg)
         assert self.execute(
             'echo deb "http://archive.ubuntu.com/ubuntu '
             '$(lsb_release -sc)-proposed main" >> '
@@ -210,11 +211,11 @@ def install_proposed_image(self):
         ).ok
         self._apt_update()
         assert self.execute(
-            "apt-get install -qy cloud-init -t=$(lsb_release -sc)-proposed"
+            f"apt-get install -qy {pkg} -t=$(lsb_release -sc)-proposed"
         ).ok
 
-    def install_ppa(self):
-        log.info("Installing PPA")
+    def install_ppa(self, pkg: str):
+        log.info("Installing %s from PPA", pkg)
         if self.execute("which add-apt-repository").failed:
             log.info("Installing missing software-properties-common package")
             self._apt_update()
@@ -223,18 +224,35 @@ def install_ppa(self):
             ).ok
         pin_origin = self.settings.CLOUD_INIT_SOURCE[4:]  # Drop leading ppa:
         pin_origin = re.sub("[^a-z0-9-]", "-", pin_origin)
+        preferences = f"""\
+package: cloud-init
+Pin: release o=LP-PPA-{pin_origin}
+Pin-Priority: 1001
+
+package: cloud-init-base
+Pin: release o=LP-PPA-{pin_origin}
+Pin-Priority: 1001
+
+package: cloud-init-cloud-sigma
+Pin: release o=LP-PPA-{pin_origin}
+Pin-Priority: 1001
+
+package: cloud-init-smart-os
+Pin: release o=LP-PPA-{pin_origin}
+Pin-Priority: 1001"""
         self.write_to_file(
             "/etc/apt/preferences.d/cloud-init-integration-testing",
-            f"package: cloud-init\nPin: release o=LP-PPA-{pin_origin}\n"
-            "Pin-Priority: 1001\n",
+            preferences,
         )
         assert self.execute(
             "add-apt-repository {} -y".format(self.settings.CLOUD_INIT_SOURCE)
         ).ok
         # PIN this PPA as priority for cloud-init installs regardless of ver
-        assert self.execute(
-            "apt-get install -qy cloud-init --allow-downgrades"
-        ).ok
+        r = self.execute(
+            f"DEBIAN_FRONTEND=noninteractive"
+            f" apt-get install -qy {pkg} --allow-downgrades"
+        )
+        assert r.ok, r.stderr
 
     @retry(tries=30, delay=1)
     def install_deb(self):
@@ -255,10 +273,10 @@ def install_deb(self):
         ).ok
 
     @retry(tries=30, delay=1)
-    def upgrade_cloud_init(self):
-        log.info("Upgrading cloud-init to latest version in archive")
+    def upgrade_cloud_init(self, pkg: str):
+        log.info("Upgrading %s to latest version in archive", pkg)
         self._apt_update()
-        assert self.execute("apt-get install -qy cloud-init").ok
+        assert self.execute(f"apt-get install -qy {pkg}").ok
 
     def _apt_update(self):
         """Run an apt update.
diff --git a/tests/integration_tests/integration_settings.py b/tests/integration_tests/integration_settings.py
index 368943c0d1c..66dcd7066f9 100644
--- a/tests/integration_tests/integration_settings.py
+++ b/tests/integration_tests/integration_settings.py
@@ -71,6 +71,10 @@
 #   A path to a valid package to be uploaded and installed
 CLOUD_INIT_SOURCE = "NONE"
 
+# cloud-init metapackage to install
+# Examples: cloud-init, cloud-init-base, cloud-init-smart-os
+CLOUD_INIT_PKG = "cloud-init"
+
 # Before an instance is torn down, we run `cloud-init collect-logs`
 # and transfer them locally. These settings specify when to collect these
 # logs and where to put them on the local filesystem