From 14c10fe060c16133ce4790ee261096589ef85179 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Fri, 24 Feb 2023 10:54:41 +0000
Subject: [PATCH] Add handling of additonal statuses to guard

Add handling of additional statuses to allow charms to set
waiting, maintenance etc. Also fix bug due to msg missing
from exception

Change-Id: Ib7dee525c6924e4d9f6fb17e0a0181642b32b9b5
---
 ops-sunbeam/ops_sunbeam/guard.py | 36 +++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/ops-sunbeam/ops_sunbeam/guard.py b/ops-sunbeam/ops_sunbeam/guard.py
index 71c8b7a2..581d4d26 100644
--- a/ops-sunbeam/ops_sunbeam/guard.py
+++ b/ops-sunbeam/ops_sunbeam/guard.py
@@ -24,6 +24,8 @@ from ops.charm import (
 )
 from ops.model import (
     BlockedStatus,
+    MaintenanceStatus,
+    WaitingStatus,
 )
 
 logger = logging.getLogger(__name__)
@@ -35,9 +37,29 @@ class GuardExceptionError(Exception):
     pass
 
 
-class BlockedExceptionError(Exception):
+class BaseStatusExceptionError(Exception):
     """Charm is blocked."""
 
+    def __init__(self, msg):
+        self.msg = msg
+        super().__init__(self.msg)
+
+
+class BlockedExceptionError(BaseStatusExceptionError):
+    """Charm is blocked."""
+
+    pass
+
+
+class MaintenanceExceptionError(BaseStatusExceptionError):
+    """Charm is performing maintenance."""
+
+    pass
+
+
+class WaitingExceptionError(BaseStatusExceptionError):
+    """Charm is waiting."""
+
     pass
 
 
@@ -79,6 +101,18 @@ def guard(
             "Charm is blocked in section '%s' due to '%s'", section, str(e)
         )
         charm.status.set(BlockedStatus(e.msg))
+    except WaitingExceptionError as e:
+        logger.warning(
+            "Charm is waiting in section '%s' due to '%s'", section, str(e)
+        )
+        charm.status.set(WaitingStatus(e.msg))
+    except MaintenanceExceptionError as e:
+        logger.warning(
+            "Charm performing maintenance in section '%s' due to '%s'",
+            section,
+            str(e),
+        )
+        charm.status.set(MaintenanceStatus(e.msg))
     except Exception as e:
         # something else went wrong
         if handle_exception: