Fixing state_db not having delete_field attribute causing a crash when DPUs in bad state

config/chassis_modules.py

@@ -6,7 +6,7 @@
 import subprocess
 import utilities_common.cli as clicommon
 from utilities_common.chassis import is_smartswitch, get_all_dpus
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 TIMEOUT_SECS = 10
 TRANSITION_TIMEOUT = timedelta(seconds=240)  # 4 minutes
@@ -27,6 +27,12 @@ def set_entry(self, table, key, entry):
         for field, value in entry.items():
             self.db.set("STATE_DB", redis_key, field, value)
 
+    def delete_field(self, table, key, field):
+        """Delete a specific field from table|key."""
+        redis_key = f"{table}|{key}"
+        client = self.db.get_redis_client("STATE_DB")
+        return client.hdel(redis_key, field)
+
 #
 # 'chassis_modules' group ('config chassis_modules ...')
 #
@@ -72,8 +78,9 @@ def set_state_transition_in_progress(db, chassis_module_name, value):
     entry = state_db.get_entry('CHASSIS_MODULE_TABLE', chassis_module_name) or {}
     entry['state_transition_in_progress'] = value
     if value == 'True':
-        entry['transition_start_time'] = datetime.utcnow().isoformat()
+        entry['transition_start_time'] = datetime.now(timezone.utc).isoformat()
     else:
+        # Remove transition_start_time from both local entry and database
         entry.pop('transition_start_time', None)
         state_db.delete_field('CHASSIS_MODULE_TABLE', chassis_module_name, 'transition_start_time')
     state_db.set_entry('CHASSIS_MODULE_TABLE', chassis_module_name, entry)
@@ -92,7 +99,14 @@ def is_transition_timed_out(db, chassis_module_name):
         start_time = datetime.fromisoformat(start_time_str)
     except ValueError:
         return False
-    return datetime.utcnow() - start_time > TRANSITION_TIMEOUT
+
+    # Use UTC everywhere for consistent comparison
+    current_time = datetime.now(timezone.utc)
+    if start_time.tzinfo is None:
+        # If stored time is naive, assume it's UTC
+        start_time = start_time.replace(tzinfo=timezone.utc)
+
+    return current_time - start_time > TRANSITION_TIMEOUT
 
 #
 # Name: check_config_module_state_with_timeout

tests/chassis_modules_test.py

@@ -1,7 +1,7 @@
 import sys
 import os
 from click.testing import CliRunner
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from config.chassis_modules import (
     set_state_transition_in_progress,
     is_transition_timed_out,
@@ -490,7 +490,7 @@ def test_shutdown_triggers_transition_in_progress(self):
             fvs = {
                 'admin_status': 'up',
                 'state_transition_in_progress': 'True',
-                'transition_start_time': datetime.utcnow().isoformat()
+                'transition_start_time': datetime.now(timezone.utc).isoformat()
             }
             db.cfgdb.set_entry('CHASSIS_MODULE', "DPU0", fvs)
 
@@ -517,7 +517,7 @@ def test_shutdown_triggers_transition_timeout(self):
             fvs = {
                 'admin_status': 'up',
                 'state_transition_in_progress': 'True',
-                'transition_start_time': (datetime.utcnow() - timedelta(minutes=30)).isoformat()
+                'transition_start_time': (datetime.now(timezone.utc) - timedelta(minutes=30)).isoformat()
             }
             db.cfgdb.set_entry('CHASSIS_MODULE', "DPU0", fvs)
 
@@ -593,15 +593,35 @@ def test_is_transition_timed_out_all_paths(self):
         assert is_transition_timed_out(db, "DPU0") is False
 
         # Case 4: Timed out
-        old_time = (datetime.utcnow() - TRANSITION_TIMEOUT - timedelta(seconds=1)).isoformat()
+        old_time = (datetime.now(timezone.utc) - TRANSITION_TIMEOUT - timedelta(seconds=1)).isoformat()
         db.statedb.get_entry.return_value = {"transition_start_time": old_time}
         assert is_transition_timed_out(db, "DPU0") is True
 
         # Case 5: Not timed out yet
-        now = datetime.utcnow().isoformat()
+        now = datetime.now(timezone.utc).isoformat()
         db.statedb.get_entry.return_value = {"transition_start_time": now}
         assert is_transition_timed_out(db, "DPU0") is False
 
+    def test_delete_field(self):
+        """Single test to cover missing delete_field and timezone handling lines"""
+        from config.chassis_modules import StateDBHelper
+        from datetime import timezone
+
+        # Test delete_field method (covers lines 32-34)
+        mock_sonic_db = mock.MagicMock()
+        mock_redis_client = mock.MagicMock()
+        mock_sonic_db.get_redis_client.return_value = mock_redis_client
+        helper = StateDBHelper(mock_sonic_db)
+        helper.delete_field('TEST_TABLE', 'test_key', 'test_field')
+        mock_redis_client.hdel.assert_called_once_with("TEST_TABLE|test_key", "test_field")
+
+        # Test timezone-aware datetime handling (covers line 109)
+        db = mock.MagicMock()
+        db.statedb = mock.MagicMock()
+        tz_time = (datetime.now(timezone.utc) - TRANSITION_TIMEOUT - timedelta(seconds=1)).isoformat()
+        db.statedb.get_entry.return_value = {"transition_start_time": tz_time}
+        assert is_transition_timed_out(db, "DPU0") is True
+
     @classmethod
     def teardown_class(cls):
         print("TEARDOWN")