bug: properly wait for removal of units and machines

anvil-python/anvil/commands/haproxy.py

@@ -26,10 +26,10 @@
 from sunbeam.jobs.steps import (
     AddMachineUnitsStep,
     DeployMachineApplicationStep,
-    RemoveMachineUnitStep,
 )
 
 from anvil.jobs.manifest import Manifest
+from anvil.jobs.steps import RemoveMachineUnitStep
 
 LOG = logging.getLogger(__name__)
 

anvil-python/anvil/commands/juju.py

@@ -13,12 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from os import environ
 import os.path
 import subprocess
 
 from rich.status import Status
+from sunbeam.commands.juju import (
+    RemoveJujuMachineStep as SunbeamRemoveJujuMachineStep,
+)
 from sunbeam.jobs.common import BaseStep, Result, ResultType
+from sunbeam.jobs.juju import CONTROLLER_MODEL
+
+LOG = logging.getLogger(__name__)
 
 
 class JujuAddSSHKeyStep(BaseStep):
@@ -54,3 +61,36 @@ def run(self, status: Status | None) -> Result:
                 message="Could not find public ssh key (~/.ssh/id_rsa.pub)",
             )
         return Result(ResultType.COMPLETED)
+
+
+class RemoveJujuMachineStep(SunbeamRemoveJujuMachineStep):
+    def run(self, status: Status | None = None) -> Result:
+        res = super().run(status)
+        if res.result_type != ResultType.COMPLETED:
+            return res
+        try:
+            cmd = [
+                self._get_juju_binary(),
+                "wait-for",
+                "machine",
+                "-m",
+                CONTROLLER_MODEL,
+                str(self.machine_id),
+                "--query",
+                'life=="dead"',
+            ]
+            LOG.debug(f'Running command {" ".join(cmd)}')
+            process = subprocess.run(
+                cmd, capture_output=True, text=True, check=True
+            )
+            LOG.debug(
+                f"Command finished. stdout={process.stdout}, stderr={process.stderr}"
+            )
+        except subprocess.CalledProcessError as e:
+            LOG.exception(
+                f"Error waiting for removal of machine {self.machine_id} from Juju"
+            )
+            LOG.warning(e.stderr)
+            return Result(ResultType.FAILED, str(e))
+
+        return Result(ResultType.COMPLETED)

anvil-python/anvil/commands/maas_agent.py

@@ -22,10 +22,10 @@
 from sunbeam.jobs.steps import (
     AddMachineUnitsStep,
     DeployMachineApplicationStep,
-    RemoveMachineUnitStep,
 )
 
 from anvil.jobs.manifest import Manifest
+from anvil.jobs.steps import RemoveMachineUnitStep
 
 APPLICATION = "maas-agent"
 CONFIG_KEY = "TerraformVarsMaasagentPlan"

anvil-python/anvil/commands/maas_region.py

@@ -22,10 +22,10 @@
 from sunbeam.jobs.steps import (
     AddMachineUnitsStep,
     DeployMachineApplicationStep,
-    RemoveMachineUnitStep,
 )
 
 from anvil.jobs.manifest import Manifest
+from anvil.jobs.steps import RemoveMachineUnitStep
 
 APPLICATION = "maas-region"
 CONFIG_KEY = "TerraformVarsMaasregionPlan"

anvil-python/anvil/commands/postgresql.py

@@ -25,10 +25,10 @@
 from sunbeam.jobs.steps import (
     AddMachineUnitsStep,
     DeployMachineApplicationStep,
-    RemoveMachineUnitStep,
 )
 
 from anvil.jobs.manifest import Manifest
+from anvil.jobs.steps import RemoveMachineUnitStep
 
 LOG = logging.getLogger(__name__)
 APPLICATION = "postgresql"

anvil-python/anvil/jobs/steps.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2024 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import subprocess
+
+from rich.status import Status
+from sunbeam.commands.juju import JujuStepHelper
+from sunbeam.jobs.common import Result, ResultType
+from sunbeam.jobs.juju import CONTROLLER_MODEL
+from sunbeam.jobs.steps import (
+    RemoveMachineUnitStep as SunbeamRemoveMachineUnitStep,
+)
+
+LOG = logging.getLogger(__name__)
+
+
+class RemoveMachineUnitStep(SunbeamRemoveMachineUnitStep, JujuStepHelper):
+    def run(self, status: Status | None = None) -> Result:
+        res = super().run(status)
+        if res.result_type != ResultType.COMPLETED:
+            return res
+        try:
+            cmd = [
+                self._get_juju_binary(),
+                "wait-for",
+                "unit",
+                "-m",
+                CONTROLLER_MODEL,
+                self.unit,
+                "--query",
+                'life=="dead"',
+            ]
+            LOG.debug(f'Running command {" ".join(cmd)}')
+            process = subprocess.run(
+                cmd, capture_output=True, text=True, check=True
+            )
+            LOG.debug(
+                f"Command finished. stdout={process.stdout}, stderr={process.stderr}"
+            )
+        except subprocess.CalledProcessError as e:
+            LOG.exception(
+                f"Error waiting for removal of unit {self.unit} from Juju"
+            )
+            LOG.warning(e.stderr)
+            return Result(ResultType.FAILED, str(e))
+
+        return Result(ResultType.COMPLETED)

anvil-python/anvil/provider/local/commands.py

@@ -42,7 +42,6 @@
     CreateJujuUserStep,
     JujuLoginStep,
     RegisterJujuUserStep,
-    RemoveJujuMachineStep,
     SaveJujuUserLocallyStep,
 )
 from sunbeam.jobs.checks import (
@@ -74,7 +73,7 @@
     RemoveHAProxyUnitStep,
     haproxy_install_steps,
 )
-from anvil.commands.juju import JujuAddSSHKeyStep
+from anvil.commands.juju import JujuAddSSHKeyStep, RemoveJujuMachineStep
 from anvil.commands.maas_agent import (
     RemoveMAASAgentUnitStep,
     maas_agent_install_steps,