From edaba9867309f1fc263869963b15c91aa1c865c5 Mon Sep 17 00:00:00 2001
From: Jack Lloyd-Walters <jack.lloyd-walters@canonical.com>
Date: Wed, 13 Nov 2024 13:33:14 +0000
Subject: [PATCH] plz?

---
 maas-region/src/charm.py  | 9 ++++-----
 maas-region/src/helper.py | 2 ++
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/maas-region/src/charm.py b/maas-region/src/charm.py
index 1f75c59..c165d53 100755
--- a/maas-region/src/charm.py
+++ b/maas-region/src/charm.py
@@ -370,9 +370,7 @@ def _on_install(self, _event: ops.InstallEvent) -> None:
 
         self._write_snap_version_()
         self._write_app_type_(self.unit, "region")
-        self._ensure_maas_cohort(_event)
-
-        _cohort = self.get_cohort()
+        _cohort = self._ensure_maas_cohort(_event) or self.get_cohort()
         if not _cohort:
             logger.exception("Snap cohort not found")
             return
@@ -589,7 +587,7 @@ def _on_maas_cluster_data_changed(self, event: ops.RelationChangedEvent) -> None
             logger.exception("Awaiting unit refresh")
             return
 
-    def _ensure_maas_cohort(self, event: ops.InstallEvent) -> None:
+    def _ensure_maas_cohort(self, event: ops.InstallEvent) -> Union[None, str]:
         logger.info(event)
         _cohort = self.get_cohort()
 
@@ -608,11 +606,12 @@ def _ensure_maas_cohort(self, event: ops.InstallEvent) -> None:
 
             self.set_cohort(_cohort)
             logger.debug(_cohort)
-            return
+            return _cohort
 
         if not _cohort:
             event.defer()
             return
+        return _cohort
 
     def _on_create_admin_action(self, event: ops.ActionEvent):
         """Handle the create-admin action.
diff --git a/maas-region/src/helper.py b/maas-region/src/helper.py
index fc362ce..3a540d7 100644
--- a/maas-region/src/helper.py
+++ b/maas-region/src/helper.py
@@ -307,6 +307,7 @@ def get_or_create_snap_cohort() -> Union[str, None]:
 
         verbose_info = maas._snap("info", ["--verbose"])
         if _found_cohort := re.search(r"cohort:\s*([^\n]+)", verbose_info):
+            logger.debug("Found cohort")
             return str(_found_cohort.group(1))
         logger.debug("Could not find cohort key in snap info")
 
@@ -314,6 +315,7 @@ def get_or_create_snap_cohort() -> Union[str, None]:
             ["sudo", "snap", "create-cohort", maas._name], universal_newlines=True
         )
         if _created_cohort := re.search(r"cohort-key:\s+([^\n]+)", cohort_creation):
+            logger.debug("Created cohort")
             return str(_created_cohort.group(1))
 
         logger.debug(f"Could not find cohort key in snap create-cohort: {_created_cohort}")