Follow-up iRMC power driver for soft reboot/poff

This is a follow-up patch of "iRMC power driver for soft reboot and
soft power off" [1] to fix a nit of exception handling code and to
make soft reboot wait logic more robust.

  [1] I8c69904063ac0a9e042f54158a20347f0c2325e1

Change-Id: I25bd413eb0f3e6c6229bf4fa053ca6ed31a5691d
Partial-Bug: #1526226

2 files changed, 22 insertions, 15 deletions

diff --git a/ironic/drivers/modules/irmc/power.py b/ironic/drivers/modules/irmc/power.py
index 368a5eede..43ed445b8 100644
--- a/ironic/drivers/modules/irmc/power.py
+++ b/ironic/drivers/modules/irmc/power.py
@@ -90,8 +90,8 @@ def _wait_power_state(task, target_state, timeout=None):
     """Wait for having changed to the target power state.
 
     :param task: A TaskManager instance containing the node to act on.
-    :raises: IRMCOperationError if the target state acknowledge failure
-         or SNMP failure.
+    :raises: IRMCOperationError if the target state acknowledge failed.
+    :raises: SNMPFailure if SNMP request failed.
     """
     node = task.node
     d_info = irmc_common.parse_driver_info(node)
@@ -176,21 +176,11 @@ def _set_power_state(task, target_state, timeout=None):
     try:
         irmc_client(STATES_MAP[target_state])
 
-        if target_state in (states.SOFT_REBOOT, states.SOFT_POWER_OFF):
-            _wait_power_state(task, target_state, timeout=timeout)
-
     except KeyError:
         msg = _("_set_power_state called with invalid power state "
                 "'%s'") % target_state
         raise exception.InvalidParameterValue(msg)
 
-    except exception.SNMPFailure as snmp_exception:
-        LOG.error(_LE("iRMC failed to acknowledge the target state "
-                      "for node %(node_id)s. Error: %(error)s"),
-                  {'node_id': node.uuid, 'error': snmp_exception})
-        raise exception.IRMCOperationError(operation=target_state,
-                                           error=snmp_exception)
-
     except scci.SCCIClientError as irmc_exception:
         LOG.error(_LE("iRMC set_power_state failed to set state to %(tstate)s "
                       " for node %(node_id)s with error: %(error)s"),
@@ -200,6 +190,22 @@ def _set_power_state(task, target_state, timeout=None):
         raise exception.IRMCOperationError(operation=operation,
                                            error=irmc_exception)
 
+    try:
+        if target_state in (states.SOFT_REBOOT, states.SOFT_POWER_OFF):
+            # note (naohirot):
+            # The following call covers both cases since SOFT_REBOOT matches
+            # 'unknown' and SOFT_POWER_OFF matches 'off' or 'unknown'.
+            _wait_power_state(task, states.SOFT_POWER_OFF, timeout=timeout)
+        if target_state == states.SOFT_REBOOT:
+            _wait_power_state(task, states.SOFT_REBOOT, timeout=timeout)
+
+    except exception.SNMPFailure as snmp_exception:
+        LOG.error(_LE("iRMC failed to acknowledge the target state "
+                      "for node %(node_id)s. Error: %(error)s"),
+                  {'node_id': node.uuid, 'error': snmp_exception})
+        raise exception.IRMCOperationError(operation=target_state,
+                                           error=snmp_exception)
+
 
 class IRMCPower(base.PowerInterface):
     """Interface for power-related actions."""
diff --git a/ironic/tests/unit/drivers/modules/irmc/test_power.py b/ironic/tests/unit/drivers/modules/irmc/test_power.py
index 72ee1d8ed..e64f616cd 100644
--- a/ironic/tests/unit/drivers/modules/irmc/test_power.py
+++ b/ironic/tests/unit/drivers/modules/irmc/test_power.py
@@ -199,8 +199,9 @@ class IRMCPowerInternalMethodsTestCase(db_base.DbTestCase):
             irmc_power._set_power_state(task, target_state)
             attach_boot_iso_if_needed_mock.assert_called_once_with(task)
         irmc_client.assert_called_once_with(irmc_power.scci.POWER_SOFT_CYCLE)
-        _wait_power_state_mock.assert_called_once_with(task, target_state,
-                                                       timeout=None)
+        _wait_power_state_mock.assert_has_calls(
+            [mock.call(task, states.SOFT_POWER_OFF, timeout=None),
+             mock.call(task, states.SOFT_REBOOT, timeout=None)])
 
     @mock.patch.object(irmc_power, '_wait_power_state', spec_set=True,
                        autospec=True)
@@ -279,7 +280,7 @@ class IRMCPowerInternalMethodsTestCase(db_base.DbTestCase):
             attach_boot_iso_if_needed_mock.assert_called_once_with(
                 task)
             _wait_power_state_mock.assert_called_once_with(
-                task, target_state, timeout=None)
+                task, states.SOFT_POWER_OFF, timeout=None)
 
 
 class IRMCPowerTestCase(db_base.DbTestCase):