Applying review comments

Author: petyaslavova

redis/maintenance_events.py

@@ -11,7 +11,7 @@
 class MaintenanceState(enum.Enum):
     NONE = "none"
     MOVING = "moving"
-    MIGRATING = "migrating"
+    MAINTENANCE = "maintenance"
     FAILING_OVER = "failing_over"
 
 
@@ -557,23 +557,32 @@ def handle_node_moved_event(self, event: NodeMovingEvent):
 
 
 class MaintenanceEventConnectionHandler:
+    # 1 = "starting maintenance" events, 0 = "completed maintenance" events
+    _EVENT_TYPES: dict[type["MaintenanceEvent"], int] = {
+        NodeMigratingEvent: 1,
+        NodeFailingOverEvent: 1,
+        NodeMigratedEvent: 0,
+        NodeFailedOverEvent: 0,
+    }
+
     def __init__(
         self, connection: "ConnectionInterface", config: MaintenanceEventsConfig
     ) -> None:
         self.connection = connection
         self.config = config
 
     def handle_event(self, event: MaintenanceEvent):
-        if isinstance(event, NodeMigratingEvent):
-            return self.handle_maintenance_start_event(MaintenanceState.MIGRATING)
-        elif isinstance(event, NodeMigratedEvent):
-            return self.handle_maintenance_completed_event()
-        elif isinstance(event, NodeFailingOverEvent):
-            return self.handle_maintenance_start_event(MaintenanceState.FAILING_OVER)
-        elif isinstance(event, NodeFailedOverEvent):
-            return self.handle_maintenance_completed_event()
-        else:
+        # get the event type by checking its class in the _EVENT_TYPES dict
+        event_type = self._EVENT_TYPES.get(event.__class__)
+
+        if event_type is None:
             logging.error(f"Unhandled event type: {event}")
+            return
+
+        if event_type:
+            self.handle_maintenance_start_event(MaintenanceState.MAINTENANCE)
+        else:
+            self.handle_maintenance_completed_event()
 
     def handle_maintenance_start_event(self, maintenance_state: MaintenanceState):
         if (

tests/test_maintenance_events.py

@@ -562,7 +562,7 @@ def test_handle_event_migrating(self):
             self.handler, "handle_maintenance_start_event"
         ) as mock_handle:
             self.handler.handle_event(event)
-            mock_handle.assert_called_once_with(MaintenanceState.MIGRATING)
+            mock_handle.assert_called_once_with(MaintenanceState.MAINTENANCE)
 
     def test_handle_event_migrated(self):
         """Test handling of NodeMigratedEvent."""
@@ -582,7 +582,7 @@ def test_handle_event_failing_over(self):
             self.handler, "handle_maintenance_start_event"
         ) as mock_handle:
             self.handler.handle_event(event)
-            mock_handle.assert_called_once_with(MaintenanceState.FAILING_OVER)
+            mock_handle.assert_called_once_with(MaintenanceState.MAINTENANCE)
 
     def test_handle_event_failed_over(self):
         """Test handling of NodeFailedOverEvent."""
@@ -608,37 +608,27 @@ def test_handle_maintenance_start_event_disabled(self):
         config = MaintenanceEventsConfig(relax_timeout=-1)
         handler = MaintenanceEventConnectionHandler(self.mock_connection, config)
 
-        result = handler.handle_maintenance_start_event(MaintenanceState.MIGRATING)
+        result = handler.handle_maintenance_start_event(MaintenanceState.MAINTENANCE)
         assert result is None
         self.mock_connection.update_current_socket_timeout.assert_not_called()
 
     def test_handle_maintenance_start_event_moving_state(self):
         """Test maintenance start event handling when connection is in MOVING state."""
         self.mock_connection.maintenance_state = MaintenanceState.MOVING
 
-        result = self.handler.handle_maintenance_start_event(MaintenanceState.MIGRATING)
+        result = self.handler.handle_maintenance_start_event(
+            MaintenanceState.MAINTENANCE
+        )
         assert result is None
         self.mock_connection.update_current_socket_timeout.assert_not_called()
 
-    def test_handle_maintenance_start_event_migrating_success(self):
+    def test_handle_maintenance_start_event_success(self):
         """Test successful maintenance start event handling for migrating."""
         self.mock_connection.maintenance_state = MaintenanceState.NONE
 
-        self.handler.handle_maintenance_start_event(MaintenanceState.MIGRATING)
-
-        assert self.mock_connection.maintenance_state == MaintenanceState.MIGRATING
-        self.mock_connection.update_current_socket_timeout.assert_called_once_with(20)
-        self.mock_connection.set_tmp_settings.assert_called_once_with(
-            tmp_relax_timeout=20
-        )
-
-    def test_handle_maintenance_start_event_failing_over_success(self):
-        """Test successful maintenance start event handling for failing over."""
-        self.mock_connection.maintenance_state = MaintenanceState.NONE
-
-        self.handler.handle_maintenance_start_event(MaintenanceState.FAILING_OVER)
+        self.handler.handle_maintenance_start_event(MaintenanceState.MAINTENANCE)
 
-        assert self.mock_connection.maintenance_state == MaintenanceState.FAILING_OVER
+        assert self.mock_connection.maintenance_state == MaintenanceState.MAINTENANCE
         self.mock_connection.update_current_socket_timeout.assert_called_once_with(20)
         self.mock_connection.set_tmp_settings.assert_called_once_with(
             tmp_relax_timeout=20
@@ -663,7 +653,7 @@ def test_handle_maintenance_completed_event_moving_state(self):
 
     def test_handle_maintenance_completed_event_success(self):
         """Test successful maintenance completed event handling."""
-        self.mock_connection.maintenance_state = MaintenanceState.MIGRATING
+        self.mock_connection.maintenance_state = MaintenanceState.MAINTENANCE
 
         self.handler.handle_maintenance_completed_event()
 

tests/test_maintenance_events_handling.py

@@ -1808,7 +1808,7 @@ def test_migrating_after_moving_multiple_proxies(self, pool_class):
         conn_event_handler = conn._maintenance_event_connection_handler
         conn_event_handler.handle_event(NodeMigratingEvent(id=3, ttl=1))
         # validate connection is in MIGRATING state
-        assert conn.maintenance_state == MaintenanceState.MIGRATING
+        assert conn.maintenance_state == MaintenanceState.MAINTENANCE
         assert conn.socket_timeout == self.config.relax_timeout
 
         # Send MIGRATED event to con with ip = key3