From 401ac34468368141dbdd521611826fd80eb54dfb Mon Sep 17 00:00:00 2001
From: Paolo Gentili <paolo.gentili@canonical.com>
Date: Thu, 1 Aug 2024 12:54:19 +0200
Subject: [PATCH 1/6] Fix: b64encoded base_user_data

---
 .../devices/zapper_kvm/README.md              |  2 +-
 .../devices/zapper_kvm/__init__.py            | 20 ++++++++++
 .../zapper_kvm/tests/test_zapper_kvm.py       | 38 ++++++++++++++++++-
 docs/reference/device-connector-types.rst     |  3 +-
 4 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
index 0329a19a..ba0f3db0 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
@@ -21,7 +21,7 @@ Unless specified via _autoinstall_ storage filter, the tool will select the larg
 - **storage_layout**: can be either `lvm`, `direct`, `zfs` or `hybrid` (Core, Desktop 23.10+)
 - **robot_tasks**: list of Zapper Robot tasks to run after a hard reset in order to follow the `autoinstall` installation
 - **cmdline_append** (optional): kernel parameters to append at the end of GRUB entry cmdline
-- **base_user_data** (optional): a custom base user-data file, it should be validated against [this schema](https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-schema.html)
+- **base_user_data** (optional): a string containing base64 encoded autoinstall user-data to use as base for provisioning, it should be validated against [this schema](https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-schema.html)
 
 ## Ubuntu Desktop OEM
 
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
index 9f74b814..447c1658 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
@@ -14,11 +14,14 @@
 
 """Zapper Connector for KVM provisioning."""
 
+import base64
+import binascii
 import logging
 import os
 import subprocess
 from typing import Any, Dict, Optional, Tuple
 
+import yaml
 from testflinger_device_connectors.devices import ProvisioningError
 from testflinger_device_connectors.devices.zapper import ZapperConnector
 from testflinger_device_connectors.devices.oemscript import OemScript
@@ -36,6 +39,22 @@ class DeviceConnector(ZapperConnector):
 
     PROVISION_METHOD = "ProvisioningKVM"
 
+    def _validate_user_data(self, encoded_user_data: str):
+        """
+        Assert `base_user_data` argument is a valid base64 encoded YAML.
+        """
+        try:
+            user_data = base64.b64decode(encoded_user_data.encode()).decode()
+            yaml.safe_load(user_data)
+        except (binascii.Error, ValueError) as exc:
+            raise ProvisioningError(
+                "Provided `base_user_data` is not base64 encoded."
+            ) from exc
+        except yaml.YAMLError as exc:
+            raise ProvisioningError(
+                "Provided `base_user_data` is not a valid YAML."
+            ) from exc
+
     def _get_autoinstall_conf(self) -> Optional[Dict[str, Any]]:
         """Prepare autoinstall-related configuration."""
         provision = self.job_data["provision_data"]
@@ -45,6 +64,7 @@ def _get_autoinstall_conf(self) -> Optional[Dict[str, Any]]:
 
         autoinstall_conf = {"storage_layout": provision["storage_layout"]}
         if "base_user_data" in provision:
+            self._validate_user_data(provision["base_user_data"])
             autoinstall_conf["base_user_data"] = provision["base_user_data"]
 
         with open(os.path.expanduser("~/.ssh/id_rsa.pub")) as pub:
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
index eae54a22..92bf1fe1 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
@@ -13,8 +13,10 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """Unit tests for Zapper KVM device connector."""
 
+import base64
 import subprocess
 import unittest
+import yaml
 from unittest.mock import Mock, patch, mock_open
 from testflinger_device_connectors.devices import ProvisioningError
 from testflinger_device_connectors.devices.zapper_kvm import DeviceConnector
@@ -242,7 +244,7 @@ def test_get_autoinstall_conf_full(self):
                     "another.robot",
                 ],
                 "storage_layout": "lvm",
-                "base_user_data": "base data content",
+                "base_user_data": "content",
             },
             "test_data": {
                 "test_username": "username",
@@ -250,16 +252,48 @@ def test_get_autoinstall_conf_full(self):
             },
         }
 
+        connector._validate_user_data = Mock()
         with patch("builtins.open", mock_open(read_data="mykey")):
             conf = connector._get_autoinstall_conf()
 
+        connector._validate_user_data.assert_called_once_with("content")
         expected = {
             "storage_layout": "lvm",
-            "base_user_data": "base data content",
+            "base_user_data": "content",
             "authorized_keys": ["mykey"],
         }
         self.assertDictEqual(conf, expected)
 
+    def test_validate_user_data(self):
+        """
+        Test whether the function returns without errors in case of a
+        sane base64 encoded YAML.
+        """
+        connector = DeviceConnector()
+        user_data_base = yaml.safe_dump({"key1": 1, "key2": [1, 2, 3]})
+        encoded = base64.b64encode(user_data_base.encode()).decode()
+        connector._validate_user_data(encoded)
+
+    def test_validate_user_data_raises_decode(self):
+        """
+        Test whether the function raises an exception if the input
+        is not correctly encoded.
+        """
+        connector = DeviceConnector()
+        with self.assertRaises(ProvisioningError):
+            connector._validate_user_data("notbase64")
+
+    def test_validate_user_data_raises_load(self):
+        """
+        Test whether the function raises an exception if the input
+        is not a valid YAML.
+        """
+        connector = DeviceConnector()
+        user_data_base = "not: a: correctly: formatted: yaml"
+        encoded = base64.b64encode(user_data_base.encode()).decode()
+        with self.assertRaises(ProvisioningError):
+            connector._validate_user_data(encoded)
+
     def test_run_oem_no_url(self):
         """
         Test the function returns without further action when URL
diff --git a/docs/reference/device-connector-types.rst b/docs/reference/device-connector-types.rst
index b8828ed0..837453ee 100644
--- a/docs/reference/device-connector-types.rst
+++ b/docs/reference/device-connector-types.rst
@@ -289,8 +289,7 @@ The ``zapper_kvm`` device connector, depending on the target image, supports the
       - When provisioning an image supporting *autoinstall*, the cmdline_append can
         be used to append Kernel parameters to the standard GRUB entry.
     * - ``base_user_data``
-      - When provisioning an image supporting *autoinstall*, the base_user_data can
-        e used to provide a base user_data file instead of the basic one hosted by Zapper.
+      - An optional string containing base64 encoded user-data to use as base for autoinstall-driven provisioning.
         For more information, see
         `Autoinstall Reference <https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-reference.html>`_.
         on this topic

From 67f282f69ea4606df68b2b4c710a4e58e339c1f4 Mon Sep 17 00:00:00 2001
From: Paolo Gentili <paolo.gentili@canonical.com>
Date: Thu, 1 Aug 2024 13:18:24 +0200
Subject: [PATCH 2/6] Add: flag for 24.04 OEM via autoinstall

---
 .../devices/zapper_kvm/README.md                      | 11 +++--------
 .../devices/zapper_kvm/__init__.py                    |  6 +++++-
 .../devices/zapper_kvm/tests/test_zapper_kvm.py       |  3 +++
 docs/reference/device-connector-types.rst             |  6 ++++--
 4 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
index ba0f3db0..37020286 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
@@ -22,10 +22,9 @@ Unless specified via _autoinstall_ storage filter, the tool will select the larg
 - **robot_tasks**: list of Zapper Robot tasks to run after a hard reset in order to follow the `autoinstall` installation
 - **cmdline_append** (optional): kernel parameters to append at the end of GRUB entry cmdline
 - **base_user_data** (optional): a string containing base64 encoded autoinstall user-data to use as base for provisioning, it should be validated against [this schema](https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-schema.html)
+- __autoinstall_oem__: Set to "true" to install OEM meta-packages and the reset partition (Desktop 24.04+)
 
-## Ubuntu Desktop OEM
-
-### Jammy
+## Ubuntu Desktop 22.04 OEM
 
 Ubuntu OEM 22.04 is a two step process:
 
@@ -38,17 +37,13 @@ The tool will select the storage device with the following priority:
 2. NVME
 3. SATA
 
-#### Job parameters
+### Job parameters
 
 - __alloem_url__: URL to the `alloem` image
 - __url__ (optional): URL to the image to test, will be installed via the OEM script
 - __password__: password to configure
 - **robot_tasks**: list of Zapper Robot tasks to run after a hard reset in order to follow the `alloem` installation
 
-### Noble
-
-Ubuntu OEM 24.04 uses `autoinstall`. The procedure and the arguments are the same as _vanilla_ Ubuntu.
-
 ## Live ISO
 
 Support for live ISOs is simply performed booting from an external storage device and returning right after KVM interactions.
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
index 447c1658..1e934b41 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
@@ -62,7 +62,11 @@ def _get_autoinstall_conf(self) -> Optional[Dict[str, Any]]:
         if "storage_layout" not in provision:
             return None
 
-        autoinstall_conf = {"storage_layout": provision["storage_layout"]}
+        autoinstall_conf = {
+            "storage_layout": provision["storage_layout"],
+            "oem": provision.get("autoinstall_oem", False),
+        }
+
         if "base_user_data" in provision:
             self._validate_user_data(provision["base_user_data"])
             autoinstall_conf["base_user_data"] = provision["base_user_data"]
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
index 92bf1fe1..ec1aceed 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
@@ -224,6 +224,7 @@ def test_get_autoinstall_conf(self):
         expected = {
             "storage_layout": "lvm",
             "authorized_keys": ["mykey"],
+            "oem": False,
         }
         self.assertDictEqual(conf, expected)
 
@@ -245,6 +246,7 @@ def test_get_autoinstall_conf_full(self):
                 ],
                 "storage_layout": "lvm",
                 "base_user_data": "content",
+                "autoinstall_oem": True,
             },
             "test_data": {
                 "test_username": "username",
@@ -261,6 +263,7 @@ def test_get_autoinstall_conf_full(self):
             "storage_layout": "lvm",
             "base_user_data": "content",
             "authorized_keys": ["mykey"],
+            "oem": True,
         }
         self.assertDictEqual(conf, expected)
 
diff --git a/docs/reference/device-connector-types.rst b/docs/reference/device-connector-types.rst
index 837453ee..8c1126da 100644
--- a/docs/reference/device-connector-types.rst
+++ b/docs/reference/device-connector-types.rst
@@ -286,13 +286,15 @@ The ``zapper_kvm`` device connector, depending on the target image, supports the
       - When provisioning an image supporting *autoinstall*, the storage_layout can
         be either ``lvm`` (default), ``direct``, ``zfs`` or ``hybrid`` (Core, Desktop 23.10+)
     * - ``cmdline_append``
-      - When provisioning an image supporting *autoinstall*, the cmdline_append can
+      - (Optional) When provisioning an image supporting *autoinstall*, the cmdline_append can
         be used to append Kernel parameters to the standard GRUB entry.
     * - ``base_user_data``
-      - An optional string containing base64 encoded user-data to use as base for autoinstall-driven provisioning.
+      - (Optional) A string containing base64 encoded user-data to use as base for autoinstall-driven provisioning.
         For more information, see
         `Autoinstall Reference <https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-reference.html>`_.
         on this topic
+    * - ``autoinstall_oem``:
+      - (Optional) Set to "true" to install OEM meta-packages and the reset partition (Desktop 24.04+).
 
 .. list-table:: Supported ``provision_data`` keys for ``zapper_kvm`` with target Ubuntu OEM 22.04
     :header-rows: 1

From 0512b273c723f851b95e8ab382a3bb3e2e98741d Mon Sep 17 00:00:00 2001
From: Paolo Gentili <paolo.gentili@canonical.com>
Date: Thu, 1 Aug 2024 14:40:52 +0200
Subject: [PATCH 3/6] Change: refactoring the zapper_kvm readme

---
 .../devices/zapper_kvm/README.md                   | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
index 37020286..295d0478 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
@@ -7,9 +7,9 @@ Zapper-driven provisioning method that makes use of KVM assertions and actions.
 
 Support for vanilla Ubuntu is provided by [autoinstall](https://canonical-subiquity.readthedocs-hosted.com/en/latest/intro-to-autoinstall.html). Supported Ubuntu versions are:
 
-- Core24 (experimental)
 - Desktop >= 23.04
 - Server >= 20.04
+- UC24 (experimental)
 
 Unless specified via _autoinstall_ storage filter, the tool will select the largest storage device on the DUT. See [supported layouts](https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-reference.html#supported-layouts) for more information.
 
@@ -18,11 +18,11 @@ Unless specified via _autoinstall_ storage filter, the tool will select the larg
 - __url__: URL to the image to install
 - __username__: username to configure
 - __password__: password to configure
-- **storage_layout**: can be either `lvm`, `direct`, `zfs` or `hybrid` (Core, Desktop 23.10+)
-- **robot_tasks**: list of Zapper Robot tasks to run after a hard reset in order to follow the `autoinstall` installation
-- **cmdline_append** (optional): kernel parameters to append at the end of GRUB entry cmdline
-- **base_user_data** (optional): a string containing base64 encoded autoinstall user-data to use as base for provisioning, it should be validated against [this schema](https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-schema.html)
-- __autoinstall_oem__: Set to "true" to install OEM meta-packages and the reset partition (Desktop 24.04+)
+- __storage_layout__: can be either `lvm`, `direct`, `zfs` or `hybrid` (Desktop 23.10+, UC24)
+- __robot_tasks__: list of Zapper Robot tasks to run after a hard reset in order to follow the `autoinstall` installation
+- __cmdline_append__ (optional): kernel parameters to append at the end of GRUB entry cmdline
+- __base_user_data__ (optional): a string containing base64 encoded autoinstall user-data to use as base for provisioning, it should be validated against [this schema](https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-schema.html)
+- __autoinstall_oem__: (optional): set to "true" to install OEM meta-packages and the reset partition (Desktop 24.04+)
 
 ## Ubuntu Desktop 22.04 OEM
 
@@ -42,7 +42,7 @@ The tool will select the storage device with the following priority:
 - __alloem_url__: URL to the `alloem` image
 - __url__ (optional): URL to the image to test, will be installed via the OEM script
 - __password__: password to configure
-- **robot_tasks**: list of Zapper Robot tasks to run after a hard reset in order to follow the `alloem` installation
+- __robot_tasks__: list of Zapper Robot tasks to run after a hard reset in order to follow the `alloem` installation
 
 ## Live ISO
 

From 9558ac75f8c788c94a89e8adbc314c4efe3e9309 Mon Sep 17 00:00:00 2001
From: Paolo Gentili <paolo.gentili@canonical.com>
Date: Thu, 1 Aug 2024 14:49:37 +0200
Subject: [PATCH 4/6] Add: ubuntu_sso_email assertion when provisioning UC via
 zapper_kvm

---
 .../devices/zapper_kvm/README.md              |  1 +
 .../devices/zapper_kvm/__init__.py            | 12 ++++++
 .../zapper_kvm/tests/test_zapper_kvm.py       | 40 +++++++++++++++++++
 docs/.wordlist.txt                            |  2 +
 docs/reference/device-connector-types.rst     |  5 ++-
 5 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
index 295d0478..238fa1ed 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/README.md
@@ -23,6 +23,7 @@ Unless specified via _autoinstall_ storage filter, the tool will select the larg
 - __cmdline_append__ (optional): kernel parameters to append at the end of GRUB entry cmdline
 - __base_user_data__ (optional): a string containing base64 encoded autoinstall user-data to use as base for provisioning, it should be validated against [this schema](https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-schema.html)
 - __autoinstall_oem__: (optional): set to "true" to install OEM meta-packages and the reset partition (Desktop 24.04+)
+- __ubuntu_sso_email__: (optional): a valid Ubuntu SSO email to which the DUT provisioned with a non-dangerous grade UC image will be linked
 
 ## Ubuntu Desktop 22.04 OEM
 
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
index 1e934b41..a763de9d 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
@@ -101,6 +101,17 @@ def _validate_configuration(
             )
             retries = self.job_data["provision_data"].get("robot_retries", 1)
 
+        # If a SSO email is specified, e.g. UC, username must match
+        # the local-part because that would be the only user available
+        # on the DUT after provisioning.
+        if "ubuntu_sso_email" in self.job_data["provision_data"]:
+            email = self.job_data["provision_data"]["ubuntu_sso_email"]
+            if username != email.split("@")[0]:
+                raise ProvisioningError(
+                    "Test username doesn't match the provided "
+                    "ubuntu_sso_email."
+                )
+
         provisioning_data = {
             "url": url,
             "username": username,
@@ -119,6 +130,7 @@ def _validate_configuration(
             "skip_download",
             "wait_until_ssh",
             "live_image",
+            "ubuntu_sso_email",
         ]
         provisioning_data.update(
             {
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
index ec1aceed..c754c8b9 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
@@ -92,6 +92,7 @@ def test_validate_configuration_w_opt(self):
                 "skip_download": True,
                 "wait_until_ssh": True,
                 "live_image": False,
+                "ubuntu_sso_email": "username@domain.com",
             },
             "test_data": {
                 "test_username": "username",
@@ -115,10 +116,49 @@ def test_validate_configuration_w_opt(self):
             "skip_download": True,
             "wait_until_ssh": True,
             "live_image": False,
+            "ubuntu_sso_email": "username@domain.com",
         }
         self.assertEqual(args, ())
         self.assertDictEqual(kwargs, expected)
 
+    def test_validate_configuration_raises_sso(self):
+        """
+        Test whether the validate_configuration raises an exception
+        if the provided test_username doesn't match with the SSO email.
+        """
+
+        connector = DeviceConnector()
+        connector.config = {
+            "device_ip": "1.1.1.1",
+            "control_host": "1.1.1.2",
+            "reboot_script": ["cmd1", "cmd2"],
+        }
+        connector.job_data = {
+            "job_queue": "queue",
+            "provision_data": {
+                "url": "http://example.com/image.iso",
+                "robot_tasks": [
+                    "job.robot",
+                    "another.robot",
+                ],
+                "storage_layout": "lvm",
+                "robot_retries": 3,
+                "cmdline_append": "more arguments",
+                "skip_download": True,
+                "wait_until_ssh": True,
+                "live_image": False,
+                "ubuntu_sso_email": "realuser@domain.com",
+            },
+            "test_data": {
+                "test_username": "username",
+                "test_password": "password",
+            },
+        }
+
+        connector._get_autoinstall_conf = Mock()
+        with self.assertRaises(ProvisioningError):
+            connector._validate_configuration()
+
     def test_validate_configuration_alloem(self):
         """
         Test whether the validate_configuration function returns
diff --git a/docs/.wordlist.txt b/docs/.wordlist.txt
index 33906e99..50a70b04 100644
--- a/docs/.wordlist.txt
+++ b/docs/.wordlist.txt
@@ -85,6 +85,7 @@ SecureBoot
 SKU
 SQA
 SSID
+SSO
 subdirectories
 subfolders
 subtree
@@ -97,6 +98,7 @@ txt
 Ubuntu
 ubuntu
 url
+UC
 UI
 URI
 USB
diff --git a/docs/reference/device-connector-types.rst b/docs/reference/device-connector-types.rst
index 8c1126da..c901d872 100644
--- a/docs/reference/device-connector-types.rst
+++ b/docs/reference/device-connector-types.rst
@@ -284,7 +284,7 @@ The ``zapper_kvm`` device connector, depending on the target image, supports the
         path from the ``robot/snippets`` path in the Zapper repository.
     * - ``storage_layout``
       - When provisioning an image supporting *autoinstall*, the storage_layout can
-        be either ``lvm`` (default), ``direct``, ``zfs`` or ``hybrid`` (Core, Desktop 23.10+)
+        be either ``lvm`` (default), ``direct``, ``zfs`` or ``hybrid`` (Desktop 23.10+, UC24)
     * - ``cmdline_append``
       - (Optional) When provisioning an image supporting *autoinstall*, the cmdline_append can
         be used to append Kernel parameters to the standard GRUB entry.
@@ -295,6 +295,9 @@ The ``zapper_kvm`` device connector, depending on the target image, supports the
         on this topic
     * - ``autoinstall_oem``:
       - (Optional) Set to "true" to install OEM meta-packages and the reset partition (Desktop 24.04+).
+    * - ``ubuntu_sso_email``:
+      - (Optional) A valid Ubuntu SSO email to which the DUT provisioned with a non-dangerous grade UC image will be linked (UC24).
+.
 
 .. list-table:: Supported ``provision_data`` keys for ``zapper_kvm`` with target Ubuntu OEM 22.04
     :header-rows: 1

From 75dc625ddcaa0f24d817fde49a2ac3646c3ee16e Mon Sep 17 00:00:00 2001
From: Paolo Gentili <paolo.gentili@canonical.com>
Date: Fri, 9 Aug 2024 09:26:05 +0200
Subject: [PATCH 5/6] Change: removed assertion over ubuntu_sso_email

---
 .../devices/zapper_kvm/__init__.py            | 11 ------
 .../zapper_kvm/tests/test_zapper_kvm.py       | 38 -------------------
 docs/reference/device-connector-types.rst     |  2 +-
 3 files changed, 1 insertion(+), 50 deletions(-)

diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
index a763de9d..76d5b0cc 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
@@ -101,17 +101,6 @@ def _validate_configuration(
             )
             retries = self.job_data["provision_data"].get("robot_retries", 1)
 
-        # If a SSO email is specified, e.g. UC, username must match
-        # the local-part because that would be the only user available
-        # on the DUT after provisioning.
-        if "ubuntu_sso_email" in self.job_data["provision_data"]:
-            email = self.job_data["provision_data"]["ubuntu_sso_email"]
-            if username != email.split("@")[0]:
-                raise ProvisioningError(
-                    "Test username doesn't match the provided "
-                    "ubuntu_sso_email."
-                )
-
         provisioning_data = {
             "url": url,
             "username": username,
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
index c754c8b9..28518c06 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
@@ -121,44 +121,6 @@ def test_validate_configuration_w_opt(self):
         self.assertEqual(args, ())
         self.assertDictEqual(kwargs, expected)
 
-    def test_validate_configuration_raises_sso(self):
-        """
-        Test whether the validate_configuration raises an exception
-        if the provided test_username doesn't match with the SSO email.
-        """
-
-        connector = DeviceConnector()
-        connector.config = {
-            "device_ip": "1.1.1.1",
-            "control_host": "1.1.1.2",
-            "reboot_script": ["cmd1", "cmd2"],
-        }
-        connector.job_data = {
-            "job_queue": "queue",
-            "provision_data": {
-                "url": "http://example.com/image.iso",
-                "robot_tasks": [
-                    "job.robot",
-                    "another.robot",
-                ],
-                "storage_layout": "lvm",
-                "robot_retries": 3,
-                "cmdline_append": "more arguments",
-                "skip_download": True,
-                "wait_until_ssh": True,
-                "live_image": False,
-                "ubuntu_sso_email": "realuser@domain.com",
-            },
-            "test_data": {
-                "test_username": "username",
-                "test_password": "password",
-            },
-        }
-
-        connector._get_autoinstall_conf = Mock()
-        with self.assertRaises(ProvisioningError):
-            connector._validate_configuration()
-
     def test_validate_configuration_alloem(self):
         """
         Test whether the validate_configuration function returns
diff --git a/docs/reference/device-connector-types.rst b/docs/reference/device-connector-types.rst
index c901d872..e4ca8e06 100644
--- a/docs/reference/device-connector-types.rst
+++ b/docs/reference/device-connector-types.rst
@@ -296,7 +296,7 @@ The ``zapper_kvm`` device connector, depending on the target image, supports the
     * - ``autoinstall_oem``:
       - (Optional) Set to "true" to install OEM meta-packages and the reset partition (Desktop 24.04+).
     * - ``ubuntu_sso_email``:
-      - (Optional) A valid Ubuntu SSO email to which the DUT provisioned with a non-dangerous grade UC image will be linked (UC24).
+      - (Optional) A valid Ubuntu SSO email to which the DUT provisioned with a non-dangerous grade UC image will be linked (UC24). Please make sure to provide the corresponding *username* in the *test_data.test_username* field.
 .
 
 .. list-table:: Supported ``provision_data`` keys for ``zapper_kvm`` with target Ubuntu OEM 22.04

From e1fda87c051bf736cb5c80f0035f6a6b7b9c6bf4 Mon Sep 17 00:00:00 2001
From: Paolo Gentili <paolo.gentili@canonical.com>
Date: Fri, 9 Aug 2024 11:08:15 +0200
Subject: [PATCH 6/6] Change: autoinstall is enable is any of the supported key
 is provided

---
 .../devices/zapper_kvm/__init__.py            | 33 ++++++++++++-------
 .../zapper_kvm/tests/test_zapper_kvm.py       | 30 ++++++++---------
 docs/reference/device-connector-types.rst     |  4 +--
 3 files changed, 38 insertions(+), 29 deletions(-)

diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
index 76d5b0cc..21c87bfb 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/__init__.py
@@ -16,6 +16,7 @@
 
 import base64
 import binascii
+import contextlib
 import logging
 import os
 import subprocess
@@ -39,7 +40,7 @@ class DeviceConnector(ZapperConnector):
 
     PROVISION_METHOD = "ProvisioningKVM"
 
-    def _validate_user_data(self, encoded_user_data: str):
+    def _validate_base_user_data(self, encoded_user_data: str):
         """
         Assert `base_user_data` argument is a valid base64 encoded YAML.
         """
@@ -56,20 +57,28 @@ def _validate_user_data(self, encoded_user_data: str):
             ) from exc
 
     def _get_autoinstall_conf(self) -> Optional[Dict[str, Any]]:
-        """Prepare autoinstall-related configuration."""
-        provision = self.job_data["provision_data"]
+        """
+        Autoinstall-related keys are pre-fixed with `autoinstall_`.
 
-        if "storage_layout" not in provision:
-            return None
+        If any of those arguments are provided and valid, the function
+        returns an autoinstall_conf dictionary, including the agent
+        SSH public key.
+        """
 
-        autoinstall_conf = {
-            "storage_layout": provision["storage_layout"],
-            "oem": provision.get("autoinstall_oem", False),
-        }
+        autoinstall_conf = {}
+        for key, value in self.job_data["provision_data"].items():
+            if "autoinstall_" not in key:
+                continue
 
-        if "base_user_data" in provision:
-            self._validate_user_data(provision["base_user_data"])
-            autoinstall_conf["base_user_data"] = provision["base_user_data"]
+            key = key.replace("autoinstall_", "")
+            with contextlib.suppress(AttributeError):
+                getattr(self, f"_validate_{key}")(value)
+
+            autoinstall_conf[key] = value
+
+        if not autoinstall_conf:
+            logger.info("Autoinstall-related keys were not provided.")
+            return None
 
         with open(os.path.expanduser("~/.ssh/id_rsa.pub")) as pub:
             autoinstall_conf["authorized_keys"] = [pub.read()]
diff --git a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
index 28518c06..82dcf2b9 100644
--- a/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
+++ b/device-connectors/src/testflinger_device_connectors/devices/zapper_kvm/tests/test_zapper_kvm.py
@@ -86,7 +86,7 @@ def test_validate_configuration_w_opt(self):
                     "job.robot",
                     "another.robot",
                 ],
-                "storage_layout": "lvm",
+                "autoinstall_storage_layout": "lvm",
                 "robot_retries": 3,
                 "cmdline_append": "more arguments",
                 "skip_download": True,
@@ -145,7 +145,7 @@ def test_validate_configuration_alloem(self):
                     "job.robot",
                     "another.robot",
                 ],
-                "storage_layout": "lvm",
+                "autoinstall_storage_layout": "lvm",
             },
             "test_data": {
                 "test_username": "username",
@@ -172,7 +172,8 @@ def test_validate_configuration_alloem(self):
     def test_get_autoinstall_none(self):
         """
         Test whether the get_autoinstall_conf function returns
-        None in case the storage_layout is not specified.
+        None in case none of the autoinstall-related keys are
+        provided.
         """
 
         connector = DeviceConnector()
@@ -212,7 +213,7 @@ def test_get_autoinstall_conf(self):
                     "job.robot",
                     "another.robot",
                 ],
-                "storage_layout": "lvm",
+                "autoinstall_storage_layout": "lvm",
             },
             "test_data": {
                 "test_username": "username",
@@ -226,7 +227,6 @@ def test_get_autoinstall_conf(self):
         expected = {
             "storage_layout": "lvm",
             "authorized_keys": ["mykey"],
-            "oem": False,
         }
         self.assertDictEqual(conf, expected)
 
@@ -246,8 +246,8 @@ def test_get_autoinstall_conf_full(self):
                     "job.robot",
                     "another.robot",
                 ],
-                "storage_layout": "lvm",
-                "base_user_data": "content",
+                "autoinstall_storage_layout": "lvm",
+                "autoinstall_base_user_data": "content",
                 "autoinstall_oem": True,
             },
             "test_data": {
@@ -256,11 +256,11 @@ def test_get_autoinstall_conf_full(self):
             },
         }
 
-        connector._validate_user_data = Mock()
+        connector._validate_base_user_data = Mock()
         with patch("builtins.open", mock_open(read_data="mykey")):
             conf = connector._get_autoinstall_conf()
 
-        connector._validate_user_data.assert_called_once_with("content")
+        connector._validate_base_user_data.assert_called_once_with("content")
         expected = {
             "storage_layout": "lvm",
             "base_user_data": "content",
@@ -269,7 +269,7 @@ def test_get_autoinstall_conf_full(self):
         }
         self.assertDictEqual(conf, expected)
 
-    def test_validate_user_data(self):
+    def test_validate_base_user_data(self):
         """
         Test whether the function returns without errors in case of a
         sane base64 encoded YAML.
@@ -277,18 +277,18 @@ def test_validate_user_data(self):
         connector = DeviceConnector()
         user_data_base = yaml.safe_dump({"key1": 1, "key2": [1, 2, 3]})
         encoded = base64.b64encode(user_data_base.encode()).decode()
-        connector._validate_user_data(encoded)
+        connector._validate_base_user_data(encoded)
 
-    def test_validate_user_data_raises_decode(self):
+    def test_validate_base_user_data_raises_decode(self):
         """
         Test whether the function raises an exception if the input
         is not correctly encoded.
         """
         connector = DeviceConnector()
         with self.assertRaises(ProvisioningError):
-            connector._validate_user_data("notbase64")
+            connector._validate_base_user_data("notbase64")
 
-    def test_validate_user_data_raises_load(self):
+    def test_validate_base_user_data_raises_load(self):
         """
         Test whether the function raises an exception if the input
         is not a valid YAML.
@@ -297,7 +297,7 @@ def test_validate_user_data_raises_load(self):
         user_data_base = "not: a: correctly: formatted: yaml"
         encoded = base64.b64encode(user_data_base.encode()).decode()
         with self.assertRaises(ProvisioningError):
-            connector._validate_user_data(encoded)
+            connector._validate_base_user_data(encoded)
 
     def test_run_oem_no_url(self):
         """
diff --git a/docs/reference/device-connector-types.rst b/docs/reference/device-connector-types.rst
index e4ca8e06..9a15184f 100644
--- a/docs/reference/device-connector-types.rst
+++ b/docs/reference/device-connector-types.rst
@@ -282,13 +282,13 @@ The ``zapper_kvm`` device connector, depending on the target image, supports the
       - List of Zapper/Robot snippets to run in sequence after the USB storage device
         is plugged into the DUT and the system restarted. The snippet ID is the relative
         path from the ``robot/snippets`` path in the Zapper repository.
-    * - ``storage_layout``
+    * - ``autoinstall_storage_layout``
       - When provisioning an image supporting *autoinstall*, the storage_layout can
         be either ``lvm`` (default), ``direct``, ``zfs`` or ``hybrid`` (Desktop 23.10+, UC24)
     * - ``cmdline_append``
       - (Optional) When provisioning an image supporting *autoinstall*, the cmdline_append can
         be used to append Kernel parameters to the standard GRUB entry.
-    * - ``base_user_data``
+    * - ``autoinstall_base_user_data``
       - (Optional) A string containing base64 encoded user-data to use as base for autoinstall-driven provisioning.
         For more information, see
         `Autoinstall Reference <https://canonical-subiquity.readthedocs-hosted.com/en/latest/reference/autoinstall-reference.html>`_.