From 82e742ea667336798edad80a81aeba8b0e678567 Mon Sep 17 00:00:00 2001
From: Chad Smith <chad.smith@canonical.com>
Date: Mon, 5 Feb 2024 14:42:28 -0700
Subject: [PATCH] feat(cloud-init): support ubuntu_pro user-data (SC-1555)

In #cloud-config user-data, cloud-init will be deprecating the
ubuntu_advantage key in favor of ubuntu_pro to align with the
current product naming and branding.

To avoid tightly coupling ubuntu-pro-client with cloud-init
releases, add support for the new 'ubuntu_pro' config key and
retain support for the deprecated keys:
 - ubuntu-advantage
 - ubuntu_advantage

Since cloud-init doesn't target Ubuntu Xenial for stable release
updates (SRUs), the ubuntu-advantage/ubuntu_advantage keys must be
supported by cloud-init and ubuntu-pro-client in Xenial unless a
backport of latest cloud-init cc_ubuntu_advantage.py module is
published to Xenial.
---
 lib/auto_attach.py                     | 10 ++++++----
 uaclient/tests/test_lib_auto_attach.py | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/lib/auto_attach.py b/lib/auto_attach.py
index 798bfe9302..6f20cccf11 100644
--- a/lib/auto_attach.py
+++ b/lib/auto_attach.py
@@ -33,6 +33,11 @@
 
 LOG = logging.getLogger("ubuntupro.lib.auto_attach")
 
+# All known cloud-config keys which provide ubuntu pro configuration directives
+CLOUD_INIT_UA_KEYS = set(
+    ["ubuntu-advantage", "ubuntu_advantage", "ubuntu_pro"]
+)
+
 try:
     import cloudinit.stages as ci_stages  # type: ignore
 except ImportError:
@@ -54,10 +59,7 @@ def check_cloudinit_userdata_for_ua_info():
     if init is None:
         return False
 
-    if init.cfg and (
-        "ubuntu_advantage" in init.cfg.keys()
-        or "ubuntu-advantage" in init.cfg.keys()
-    ):
+    if init.cfg and CLOUD_INIT_UA_KEYS.intersection(init.cfg):
         return True
 
     return False
diff --git a/uaclient/tests/test_lib_auto_attach.py b/uaclient/tests/test_lib_auto_attach.py
index c529bc70c4..eed9107458 100644
--- a/uaclient/tests/test_lib_auto_attach.py
+++ b/uaclient/tests/test_lib_auto_attach.py
@@ -41,6 +41,20 @@ def test_check_cloudinit_data_returns_false_if_no_cloudinit(
                     "cloud_config_modules": ["ubuntu-advantage", "test"],
                 },
             ),
+            (
+                True,
+                {
+                    "ubuntu_pro": {"token": "TOKEN"},
+                    "cloud_config_modules": ["ubuntu-advantage", "test"],
+                },
+            ),
+            (
+                True,
+                {
+                    "ubuntu-advantage": {"token": "TOKEN"},
+                    "cloud_config_modules": ["ubuntu-advantage", "test"],
+                },
+            ),
         ),
     )
     @mock.patch("lib.auto_attach.get_cloudinit_init_stage")