Misc fixes after review feedback

sunbeam-python/sunbeam/commands/juju.py

@@ -22,7 +22,7 @@
 import subprocess
 import tempfile
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
 import pexpect
 import pwgen
@@ -172,29 +172,8 @@ def add_cloud(self, cloud_type: str, cloud_name: str) -> bool:
 
         return True
 
-    def get_available_charm_revision(
-        self, charm_name: str, track: str, risk: str, arch: str = "amd64"
-    ) -> int:
-        """Find the latest available revision of a charm in a given channel
-
-        :param charm_name: Name of charm to look up
-        :param track: Track of charm
-        :param risk: Risk of charm
-        :param arch: Arch of charm
-        """
-        available_charm_data = self._juju_cmd(*["info", charm_name])
-        channel_data = [
-            d
-            for d in available_charm_data["channels"][track][risk]
-            if arch in d["architectures"]
-        ]
-        assert len(channel_data) == 1, "Unexpected candidate charms {}".format(
-            len(channel_data)
-        )
-        return int(channel_data[0]["revision"])
-
     def revision_update_needed(
-        self, application_name: str, model: str, status: Union[dict, None] = None
+        self, application_name: str, model: str, status: dict | None = None
     ) -> bool:
         """Check if a revision update is available for an applicaton.
 
@@ -212,15 +191,14 @@ def revision_update_needed(
         deployed_revision = int(self._extract_charm_revision(app_status["charm"]))
         charm_name = self._extract_charm_name(app_status["charm"])
         deployed_channel = self.normalise_channel(app_status["charm-channel"])
-        track = deployed_channel.split("/")[0]
-        risk = deployed_channel.split("/")[1]
-        try:
-            deployed_channel.split("/")[2]
+        if len(deployed_channel.split("/")) > 2:
             LOG.debug(f"Cannot calculate upgrade for {application_name}, branch in use")
             return False
-        except IndexError:
-            pass
-        available_revision = self.get_available_charm_revision(charm_name, track, risk)
+        available_revision = run_sync(
+            self.jhelper.get_available_charm_revision(
+                model, charm_name, deployed_channel
+            )
+        )
         return bool(available_revision > deployed_revision)
 
     def normalise_channel(self, channel: str) -> str:

sunbeam-python/sunbeam/commands/upgrades/base.py

@@ -55,7 +55,7 @@ def run(self, status: Optional[Status] = None) -> Result:
 
 class UpgradeCoordinator:
     def __init__(
-        self, jhelper: JujuHelper, tfhelper: TerraformHelper, channel: str = None
+        self, jhelper: JujuHelper, tfhelper: TerraformHelper, channel: str | None = None
     ):
         """Upgrade coordinator.
 

sunbeam-python/sunbeam/commands/upgrades/inter_channel.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import logging
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, TypedDict, Union
 
 from rich.console import Console
 from rich.status import Status
@@ -32,7 +32,7 @@
     run_plan,
     update_config,
 )
-from sunbeam.jobs.juju import JujuHelper, run_sync
+from sunbeam.jobs.juju import ChannelUpdate, JujuHelper, run_sync
 from sunbeam.versions import (
     CHARM_VERSIONS,
     MACHINE_SERVICES,
@@ -45,6 +45,28 @@
 console = Console()
 
 
+class UpgradeStrategy(TypedDict):
+    """A strategy for upgrading applications.
+
+    The strategy is a list of dicts. Each dict consists of:
+        {
+            "upgrade_f": f,
+            "applications": [apps]
+        },
+
+    upgrade_f is the function to be applied to each application (in
+    parallel) to perform the upgrade.
+
+    applications is a list of applications that can be upgraded in parallel.
+
+    Currently only apps that are upgraded with the same function can be
+    grouped together.
+    """
+
+    upgrade_f: Callable[[list[str], str], None]
+    applications: list[str]
+
+
 class BaseUpgrade(BaseStep, JujuStepHelper):
     def __init__(self, name, description, jhelper, tfhelper, model):
         """Create instance of BaseUpgrade class.
@@ -58,29 +80,15 @@ def __init__(self, name, description, jhelper, tfhelper, model):
         self.tfhelper = tfhelper
         self.model = model
 
-    def get_upgrade_strategy(self) -> List[Dict[str, Union[Callable, List]]]:
-        """Return a strategy for performing the upgrade.
-
-        The strategy is a list of dicts. Each dict consists of:
-            {
-                "upgrade_f": f,
-                "applications": [apps]
-            },
+    def get_upgrade_strategy_steps(self) -> UpgradeStrategy:
+        """Return a strategy for performing the upgrade."""
 
-        upgrade_f is the function to be applied to each application (in
-        parallel) to perform the upgrade.
-
-        applications is a list of applications that can be upgraded in parallel.
-
-        Currently only apps that are upgraded with the same function can be
-        grouped together.
-        """
         raise NotImplementedError
 
     def run(self, status: Optional[Status] = None) -> Result:
         """Run control plane and machine charm upgrade."""
         self.pre_upgrade_tasks()
-        for step in self.get_upgrade_strategy():
+        for step in self.get_upgrade_strategy_steps():
             step["upgrade_f"](step["applications"], self.model)
         self.post_upgrade_tasks()
         return Result(ResultType.COMPLETED)
@@ -97,7 +105,7 @@ def upgrade_applications(
         self,
         application_list: List[str],
         model: str,
-        expect_wls: Optional[Dict[str, str]] = None,
+        expect_wls: Optional[Dict[str, list[str]]] = None,
     ) -> None:
         """Upgrade applications.
 
@@ -112,10 +120,10 @@ def upgrade_applications(
             new_channel = self.get_new_channel(app_name, model)
             if new_channel:
                 LOG.debug(f"Upgrade needed for {app_name}")
-                batch[app_name] = {
-                    "channel": new_channel,
-                    "expected_status": expect_wls,
-                }
+                batch[app_name] = ChannelUpdate(
+                    channel=new_channel,
+                    expected_status=expect_wls,
+                )
             else:
                 LOG.debug(f"{app_name} no channel upgrade needed")
         run_sync(self.jhelper.update_applications_channel(model, batch))
@@ -174,20 +182,20 @@ def __init__(self, jhelper, tfhelper, model):
             model,
         )
 
-    def get_upgrade_strategy(self) -> List[Dict[str, Union[Callable, List]]]:
+    def get_upgrade_strategy_steps(self) -> List[Dict[str, Union[Callable, List]]]:
         """Return a strategy for performing the upgrade.
 
         Upgrade all control plane applications in parallel.
         """
-        upgrade_strategy = [
-            {
-                "upgrade_f": self.upgrade_applications,
-                "applications": list(MISC_SERVICES_K8S.keys())
+        upgrade_strategy_steps = [
+            UpgradeStrategy(
+                upgrade_f=self.upgrade_applications,
+                applications=list(MISC_SERVICES_K8S.keys())
                 + list(OVN_SERVICES_K8S.keys())  # noqa
                 + list(OPENSTACK_SERVICES_K8S.keys()),  # noqa
-            },
+            ),
         ]
-        return upgrade_strategy
+        return upgrade_strategy_steps
 
     def post_upgrade_tasks(self) -> None:
         """Update channels in terraform vars db."""
@@ -224,18 +232,17 @@ def __init__(self, jhelper, tfhelper, model):
             model,
         )
 
-    def get_upgrade_strategy(self) -> List[Dict[str, Union[Callable, List]]]:
+    def get_upgrade_strategy_steps(self) -> List[Dict[str, Union[Callable, List]]]:
         """Return a strategy for performing the upgrade.
 
         Upgrade all machine applications in parallel.
         """
-        upgrade_strategy = [
-            {
-                "upgrade_f": self.upgrade_applications,
-                "applications": MACHINE_SERVICES,
-            },
+        upgrade_strategy_steps = [
+            UpgradeStrategy(
+                upgrade_f=self.upgrade_applications, applications=MACHINE_SERVICES
+            ),
         ]
-        return upgrade_strategy
+        return upgrade_strategy_steps
 
     def post_upgrade_tasks(self) -> None:
         """Update channels in terraform vars db."""

sunbeam-python/sunbeam/jobs/juju.py

@@ -21,12 +21,13 @@
 from datetime import datetime
 from functools import wraps
 from pathlib import Path
-from typing import Awaitable, Dict, List, Optional, TypeVar, cast
+from typing import Awaitable, Dict, List, Optional, TypedDict, TypeVar, cast
 
 import pytz
 import yaml
 from juju import utils as juju_utils
 from juju.application import Application
+from juju.charmhub import CharmHub
 from juju.client import client as jujuClient
 from juju.controller import Controller
 from juju.errors import (
@@ -138,6 +139,20 @@ class UnsupportedKubeconfigException(JujuException):
     pass
 
 
+class ChannelUpdate(TypedDict):
+    """Channel Update step.
+
+    Defines a channel that needs updating to and the expected
+    state of the charm afterwards.
+
+    channel: Channel to upgrade to
+    expected_status: map of accepted statuses for "workload" and "agent"
+    """
+
+    channel: str
+    expected_status: Dict[str, List[str]]
+
+
 @dataclass
 class JujuAccount:
     user: str
@@ -721,7 +736,7 @@ async def set_application_config(self, model: str, app: str, config: dict):
     async def update_applications_channel(
         self,
         model: str,
-        updates: Dict,
+        updates: Dict[str, ChannelUpdate],
         timeout: Optional[int] = None,
     ):
         """Upgrade charm to new channel
@@ -750,11 +765,6 @@ def condition() -> bool:
                 )
                 for unit in _app.units:
                     statuses[unit.entity_id] = bool(unit.agent_status_since > timestamp)
-                    if not unit.agent_status_since > timestamp:
-                        LOG.debug(
-                            f"Waiting on {unit.entity_id}: Upgrade start {timestamp} "
-                            f"WLS update {unit.agent_status_since}"
-                        )
             return all(statuses.values())
 
         try:
@@ -797,3 +807,18 @@ async def charm_refresh(self, application_name: str, model: str):
         """
         app = await self.get_application(application_name, model)
         await app.refresh()
+
+    @controller
+    async def get_available_charm_revision(
+        self, model: str, charm_name: str, channel: str
+    ) -> int:
+        """Find the latest available revision of a charm in a given channel
+
+        :param model: Name of model
+        :param charm_name: Name of charm to look up
+        :param channel: Channel to lookup charm in
+        """
+        model_impl = await self.get_model(model)
+        available_charm_data = await CharmHub(model_impl).info(charm_name, channel)
+        version = available_charm_data["channel-map"][channel]["revision"]["version"]
+        return int(version)

sunbeam-python/sunbeam/plugins/dns/plugin.py

@@ -26,6 +26,7 @@
 from sunbeam.jobs.common import run_plan
 from sunbeam.jobs.juju import JujuHelper, run_sync
 from sunbeam.plugins.interface.v1.openstack import (
+    ApplicationChannelData,
     EnableOpenStackApplicationStep,
     OpenStackControlPlanePlugin,
     TerraformPlanLocation,
@@ -176,14 +177,14 @@ def commands(self) -> dict:
     @property
     def k8s_application_data(self):
         return {
-            "designate": {
-                "channel": OPENSTACK_CHANNEL,
-                "tfvars_channel_var": None,
-            },
-            "bind": {
-                "channel": "9/edge",
-                "tfvars_channel_var": None,
-            },
+            "designate": ApplicationChannelData(
+                channel=OPENSTACK_CHANNEL,
+                tfvars_channel_var=None,
+            ),
+            "bind": ApplicationChannelData(
+                channel="9/edge",
+                tfvars_channel_var=None,
+            ),
         }
 
     @property