Addressed review comments related to refactoring

Author: rameshraghupathy

scripts/gnoi_shutdown_daemon.py

@@ -218,8 +218,9 @@ def main():
                 time.sleep(1)
                 continue
 
-            if entry.get("state_transition_in_progress", "False") == "True" and entry.get("transition_type") == "shutdown":
-                logger.log_info(f"Shutdown request detected for {dpu_name}. Initiating gNOI reboot.")
+            type = entry.get("transition_type")
+            if entry.get("state_transition_in_progress", "False") == "True" and (type == "shutdown" or type == "reboot"):
+                logger.log_info(f"{type} request detected for {dpu_name}. Initiating gNOI reboot.")
                 try:
                     dpu_ip = get_dpu_ip(dpu_name)
                     port = get_gnmi_port(dpu_name)
@@ -276,18 +277,20 @@ def main():
                     time.sleep(STATUS_POLL_INTERVAL_SEC)
 
                 if reboot_successful:
+                    if type == "reboot":
+                        success = module_base.clear_module_state_transition(db, dpu_name)
+                        if success:
+                            logger.log_info(f"Cleared transition for {dpu_name}")
+                        else:
+                            logger.log_warning(f"Failed to clear transition for {dpu_name}")
                     logger.log_info(f"Halting the services on DPU is successful for {dpu_name}.")
                 else:
                     logger.log_warning(f"Status polling of halting the services on DPU timed out for {dpu_name}.")
 
                 # NOTE:
-                # Do NOT clear CHASSIS_MODULE_TABLE transition flags here.
-                # Per HLD and platform flow, the transition is cleared by the
-                # platform's module.py AFTER set_admin_state(down) has completed
-                # (i.e., after the module is actually taken down). This avoids
-                # prematurely unblocking other components before shutdown finishes.
-
-        time.sleep(1)
+                # The CHASSIS_MODULE_TABLE transition flag is cleared for startup/shutdown in
+                # module_base.py. The daemon does not clear it. For reboot transitions, the
+                # daemon relies on the TimeoutEnforcer thread to clear any stuck transitions.
 
 if __name__ == "__main__":
     main()

tests/gnoi_shutdown_daemon_test.py

@@ -172,6 +172,7 @@ def __init__(self, *a, **k):  # allow construction if the code instantiates Modu
                 pass
             # Support both instance and class access
             get_module_state_transition = staticmethod(_fake_transition)
+            clear_module_state_transition = staticmethod(lambda db, name: True)
 
         with patch("gnoi_shutdown_daemon.SonicV2Connector") as mock_sonic, \
              patch("gnoi_shutdown_daemon.ModuleBase", new=_MBStub), \
@@ -220,6 +221,7 @@ def __init__(self, *a, **k):  # allow construction if the code instantiates Modu
             self.assertTrue(status_args[status_args.index("-rpc") + 1].endswith("RebootStatus"))
 
             all_logs = " | ".join(str(c) for c in mock_logger.method_calls)
+            self.assertIn("shutdown request detected for DPU0", all_logs)
             self.assertIn("Halting the services on DPU is successful for DPU0", all_logs)
 
 
@@ -550,6 +552,10 @@ def __init__(self, *a, **k):
     @staticmethod
     def get_module_state_transition(*_a, **_k):
         return {"state_transition_in_progress": "True", "transition_type": "shutdown"}
+    
+    @staticmethod
+    def clear_module_state_transition(db, name):
+        return True
 
 
 def _mk_pubsub_once2():
@@ -648,3 +654,141 @@ def test_shutdown_reboot_nonzero_does_not_poll_status(self):
             self.assertEqual(mock_exec.call_count, 1)
             self.assertTrue(any("reboot command failed" in str(c.args[0]).lower()
                     for c in (mock_logger.log_error.call_args_list or [])))
+
+    def test_reboot_transition_type_success(self):
+        """Test that reboot transition type is handled correctly and clears transition on success"""
+        
+        class _MBStubReboot:
+            def __init__(self, *a, **k):
+                pass
+                
+            @staticmethod
+            def get_module_state_transition(*_a, **_k):
+                return {"state_transition_in_progress": "True", "transition_type": "reboot"}
+            
+            @staticmethod
+            def clear_module_state_transition(db, name):
+                return True
+        
+        with patch("gnoi_shutdown_daemon.SonicV2Connector") as mock_sonic, \
+             patch("gnoi_shutdown_daemon.ModuleBase", new=_MBStubReboot), \
+             patch("gnoi_shutdown_daemon.execute_gnoi_command") as mock_exec, \
+             patch("gnoi_shutdown_daemon.is_tcp_open", return_value=True), \
+             patch("gnoi_shutdown_daemon._cfg_get_entry",
+                   side_effect=lambda table, key:
+                       {"ips@": "10.0.0.1"} if table == "DHCP_SERVER_IPV4_PORT" else {"gnmi_port": "8080"}), \
+             patch("gnoi_shutdown_daemon.time.sleep", return_value=None), \
+             patch("gnoi_shutdown_daemon.logger") as mock_logger:
+            import gnoi_shutdown_daemon as d
+            db = MagicMock()
+            pubsub = MagicMock()
+            pubsub.get_message.side_effect = [
+                {"type": "pmessage", "channel": "__keyspace@6__:CHASSIS_MODULE_TABLE|DPU0", "data": "set"},
+                Exception("stop"),
+            ]
+            db.pubsub.return_value = pubsub
+            mock_sonic.return_value = db
+
+            mock_exec.side_effect = [
+                (0, "OK", ""),  # Reboot command
+                (0, "reboot complete", ""),  # RebootStatus
+            ]
+
+            try:
+                d.main()
+            except Exception:
+                pass
+
+            # Should make both Reboot and RebootStatus calls
+            self.assertEqual(mock_exec.call_count, 2)
+            
+            # Check logs for reboot-specific messages
+            all_logs = " | ".join(str(c) for c in mock_logger.method_calls)
+            self.assertIn("reboot request detected for DPU0", all_logs)
+            self.assertIn("Cleared transition for DPU0", all_logs)
+            self.assertIn("Halting the services on DPU is successful for DPU0", all_logs)
+
+    def test_reboot_transition_clear_failure(self):
+        """Test that reboot transition logs warning when clear fails"""
+        
+        class _MBStubRebootFail:
+            def __init__(self, *a, **k):
+                pass
+                
+            @staticmethod
+            def get_module_state_transition(*_a, **_k):
+                return {"state_transition_in_progress": "True", "transition_type": "reboot"}
+            
+            @staticmethod
+            def clear_module_state_transition(db, name):
+                return False  # Simulate failure
+        
+        with patch("gnoi_shutdown_daemon.SonicV2Connector") as mock_sonic, \
+             patch("gnoi_shutdown_daemon.ModuleBase", new=_MBStubRebootFail), \
+             patch("gnoi_shutdown_daemon.execute_gnoi_command") as mock_exec, \
+             patch("gnoi_shutdown_daemon.is_tcp_open", return_value=True), \
+             patch("gnoi_shutdown_daemon._cfg_get_entry",
+                   side_effect=lambda table, key:
+                       {"ips@": "10.0.0.1"} if table == "DHCP_SERVER_IPV4_PORT" else {"gnmi_port": "8080"}), \
+             patch("gnoi_shutdown_daemon.time.sleep", return_value=None), \
+             patch("gnoi_shutdown_daemon.logger") as mock_logger:
+            import gnoi_shutdown_daemon as d
+            db = MagicMock()
+            pubsub = MagicMock()
+            pubsub.get_message.side_effect = [
+                {"type": "pmessage", "channel": "__keyspace@6__:CHASSIS_MODULE_TABLE|DPU0", "data": "set"},
+                Exception("stop"),
+            ]
+            db.pubsub.return_value = pubsub
+            mock_sonic.return_value = db
+
+            mock_exec.side_effect = [
+                (0, "OK", ""),  # Reboot command
+                (0, "reboot complete", ""),  # RebootStatus
+            ]
+
+            try:
+                d.main()
+            except Exception:
+                pass
+
+            # Check for warning log when clear fails
+            all_logs = " | ".join(str(c) for c in mock_logger.method_calls)
+            self.assertIn("Failed to clear transition for DPU0", all_logs)
+
+    def test_status_polling_timeout_warning(self):
+        """Test that timeout during status polling logs the appropriate warning"""
+        
+        with patch("gnoi_shutdown_daemon.SonicV2Connector") as mock_sonic, \
+             patch("gnoi_shutdown_daemon.ModuleBase", new=_MBStub2), \
+             patch("gnoi_shutdown_daemon.execute_gnoi_command") as mock_exec, \
+             patch("gnoi_shutdown_daemon.is_tcp_open", return_value=True), \
+             patch("gnoi_shutdown_daemon._cfg_get_entry",
+                   side_effect=lambda table, key:
+                       {"ips@": "10.0.0.1"} if table == "DHCP_SERVER_IPV4_PORT" else {"gnmi_port": "8080"}), \
+             patch("gnoi_shutdown_daemon.time.sleep", return_value=None), \
+             patch("gnoi_shutdown_daemon.time.monotonic", side_effect=[0, 100]), \
+             patch("gnoi_shutdown_daemon.logger") as mock_logger:
+            import gnoi_shutdown_daemon as d
+            db = MagicMock()
+            pubsub = MagicMock()
+            pubsub.get_message.side_effect = [
+                {"type": "pmessage", "channel": "__keyspace@6__:CHASSIS_MODULE_TABLE|DPU0", "data": "set"},
+                Exception("stop"),
+            ]
+            db.pubsub.return_value = pubsub
+            mock_sonic.return_value = db
+
+            mock_exec.side_effect = [
+                (0, "OK", ""),  # Reboot command
+                (0, "not complete", ""),  # RebootStatus - never returns complete
+            ]
+
+            try:
+                d.main()
+            except Exception:
+                pass
+
+            # Check for timeout warning
+            all_logs = " | ".join(str(c) for c in mock_logger.method_calls)
+            self.assertIn("Status polling of halting the services on DPU timed out for DPU0", all_logs)