From 8cf92b32d5cbe0faa8bab57a42ad6c811f9672e9 Mon Sep 17 00:00:00 2001
From: Dima Tisnek <dima.tisnek@canonical.com>
Date: Fri, 22 Nov 2024 13:00:57 +0900
Subject: [PATCH] chore: use StatusStr for known status values

---
 juju/status.py | 62 ++++++++++++++++++++++++++++++++++++--------------
 pyproject.toml |  7 +++---
 setup.py       |  1 +
 3 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/juju/status.py b/juju/status.py
index 8825f61f..624cda21 100644
--- a/juju/status.py
+++ b/juju/status.py
@@ -3,14 +3,50 @@
 from __future__ import annotations
 
 import logging
+import sys
 import warnings
 
+if sys.version_info >= (3, 11):
+    from enum import StrEnum
+else:
+    from backports.strenum import StrEnum
+
 from .client import client
 
 log = logging.getLogger(__name__)
 
 
-def derive_status(statuses: list[str]):
+class StatusStr(StrEnum):
+    """Recognised status values.
+
+    Please keep this set exact same as the severity map below.
+    """
+
+    error = "error"
+    blocked = "blocked"
+    waiting = "waiting"
+    maintenance = "maintenance"
+    active = "active"
+    terminated = "terminated"
+    unknown = "unknown"
+
+
+""" severity_map holds status values with a severity measure.
+Status values with higher severity are used in preference to others.
+"""
+severity_map: dict[StatusStr, int] = {
+    # FIXME: Juju defines a lot more status values #1204
+    StatusStr.error: 100,
+    StatusStr.blocked: 90,
+    StatusStr.waiting: 80,
+    StatusStr.maintenance: 70,
+    StatusStr.active: 60,
+    StatusStr.terminated: 50,
+    StatusStr.unknown: 40,
+}
+
+
+def derive_status(statuses: list[str | StatusStr]) -> StatusStr:
     """Derive status from a set.
 
     derive_status is used to determine the application status from a set of unit
@@ -18,27 +54,19 @@ def derive_status(statuses: list[str]):
 
     :param statuses: list of known unit workload statuses
     """
-    current = "unknown"
+    current: StatusStr = StatusStr.unknown
     for status in statuses:
-        if status in severity_map and severity_map[status] > severity_map[current]:
+        try:
+            status = StatusStr(status)
+        except ValueError:
+            continue
+        if (new_level := severity_map.get(status)) and new_level > severity_map[
+            current
+        ]:
             current = status
     return current
 
 
-""" severity_map holds status values with a severity measure.
-Status values with higher severity are used in preference to others.
-"""
-severity_map = {
-    "error": 100,
-    "blocked": 90,
-    "waiting": 80,
-    "maintenance": 70,
-    "active": 60,
-    "terminated": 50,
-    "unknown": 40,
-}
-
-
 async def formatted_status(model, target=None, raw=False, filters=None):
     """Returns a string that mimics the content of the information
     returned in the juju status command. If the raw parameter is
diff --git a/pyproject.toml b/pyproject.toml
index 6974f630..47c22809 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ description = "Python library for Juju"
 readme = "docs/readme.rst"
 license = { file = "LICENSE" }
 maintainers = [{name = "Juju Ecosystem Engineering", email = "juju@lists.ubuntu.com"}]
-requires-python = ">=3.8"
+requires-python = ">=3.8.6"
 classifiers = [
     "Development Status :: 5 - Production/Stable",
     "Intended Audience :: Developers",
@@ -33,7 +33,8 @@ dependencies = [
     "kubernetes>=12.0.1,<31.0.0",
     "hvac",
     "packaging",
-    "typing-extensions>=4.5.0"
+    "typing-extensions>=4.5.0",
+    "backports.strenum>=1.3.1",
 ]
 
 [project.urls]
@@ -212,7 +213,7 @@ ignore = [
 [tool.pyright]
 # These are tentative
 include = ["**/*.py"]
-pythonVersion = "3.8"
+pythonVersion = "3.8.6"
 typeCheckingMode = "strict"
 useLibraryCodeForTypes = true
 reportGeneralTypeIssues = true
diff --git a/setup.py b/setup.py
index d3464ec6..5b845661 100644
--- a/setup.py
+++ b/setup.py
@@ -32,6 +32,7 @@
         "hvac",
         "packaging",
         "typing-extensions>=4.5.0",
+        "backports.strenum",
     ],
     include_package_data=True,
     maintainer="Juju Ecosystem Engineering",