From a88cd8e3415cb7b56143365213eda6f8e2f16e74 Mon Sep 17 00:00:00 2001
From: srunde3 <95937393+srunde3@users.noreply.github.com>
Date: Thu, 28 Mar 2024 07:04:08 -0700
Subject: [PATCH] fix snap: return empty ubuntu pro info for non-core snap to
 fix registration bug (#234)

---
 .../manager/tests/test_ubuntuproinfo.py       | 28 +++++++++++++++++++
 landscape/client/manager/ubuntuproinfo.py     |  6 ++++
 2 files changed, 34 insertions(+)

diff --git a/landscape/client/manager/tests/test_ubuntuproinfo.py b/landscape/client/manager/tests/test_ubuntuproinfo.py
index 9d6872be6..b5f99f3cd 100644
--- a/landscape/client/manager/tests/test_ubuntuproinfo.py
+++ b/landscape/client/manager/tests/test_ubuntuproinfo.py
@@ -160,3 +160,31 @@ def test_persistence_reset(self):
         self.assertEqual(2, len(messages))
         self.assertTrue("ubuntu-pro-info" in messages[1])
         self.assertEqual(messages[1]["ubuntu-pro-info"], data)
+
+    @mock.patch("landscape.client.manager.ubuntuproinfo.IS_SNAP", new=True)
+    def test_pro_client_not_called_for_snap(self):
+        """
+        The snap will not currently allow calls to the pro client.
+
+        Ensure that get_ubuntu_pro_info returns an empty dictionary instead of
+        calling the subprocess for pro.
+        """
+        ubuntu_pro_info = get_ubuntu_pro_info()
+        self.assertEqual({}, ubuntu_pro_info)
+
+    def test_mock_info_sent_for_core_snap(self):
+        """
+        Ensure that a Core snap still receives mocked ubuntu pro info even if
+        the snap generally doesn't support *real* ubuntu pro info
+        """
+        with mock.patch.multiple(
+            "landscape.client.manager.ubuntuproinfo",
+            IS_CORE=True,
+            IS_SNAP=True,
+        ):
+            ubuntu_pro_info = get_ubuntu_pro_info()
+
+        self.assertIn("effective", ubuntu_pro_info)
+        self.assertIn("expires", ubuntu_pro_info)
+        contract = ubuntu_pro_info["contract"]
+        self.assertIn("landscape", contract["products"])
diff --git a/landscape/client/manager/ubuntuproinfo.py b/landscape/client/manager/ubuntuproinfo.py
index b582e5d45..e5c054a61 100644
--- a/landscape/client/manager/ubuntuproinfo.py
+++ b/landscape/client/manager/ubuntuproinfo.py
@@ -6,6 +6,7 @@
 from pathlib import Path
 
 from landscape.client import IS_CORE
+from landscape.client import IS_SNAP
 from landscape.client.manager.plugin import ManagerPlugin
 from landscape.lib.persist import Persist
 
@@ -78,6 +79,11 @@ def get_ubuntu_pro_info() -> dict:
             expiration_datetime,
         )
 
+    if IS_SNAP:
+        # Snap does not support Ubuntu Pro Info and throws an error if `pro` is
+        # called.
+        return {}
+
     try:
         completed_process = subprocess.run(
             ["pro", "status", "--format", "json"],