Applying review comments

Author: petyaslavova

redis/_parsers/base.py

@@ -209,6 +209,7 @@ def handle_push_response(self, response, **kwargs):
                 and self.invalidation_push_handler_func
             ):
                 return self.invalidation_push_handler_func(response)
+
             if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
                 # Expected message format is: MOVING <seq_number> <time> <endpoint>
                 id = response[1]
@@ -280,6 +281,7 @@ async def handle_push_response(self, response, **kwargs):
                 and self.invalidation_push_handler_func
             ):
                 return await self.invalidation_push_handler_func(response)
+
             if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
                 # push notification from enterprise cluster for node moving
                 id = response[1]

redis/connection.py

@@ -431,15 +431,13 @@ def __init__(
         self.maintenance_events_config = maintenance_events_config
 
         # Set up maintenance events if enabled
-        if maintenance_events_config and maintenance_events_config.enabled:
-            self._enable_maintenance_events(
-                maintenance_events_pool_handler,
-                orig_host_address,
-                orig_socket_timeout,
-                orig_socket_connect_timeout,
-            )
-        else:
-            self._maintenance_event_connection_handler = None
+        self._configure_maintenance_events(
+            maintenance_events_pool_handler,
+            orig_host_address,
+            orig_socket_timeout,
+            orig_socket_connect_timeout,
+        )
+
         self._should_reconnect = False
         self.maintenance_state = maintenance_state
 
@@ -498,7 +496,7 @@ def set_parser(self, parser_class):
         """
         self._parser = parser_class(socket_read_size=self._socket_read_size)
 
-    def _enable_maintenance_events(
+    def _configure_maintenance_events(
         self,
         maintenance_events_pool_handler=None,
         orig_host_address=None,
@@ -510,6 +508,7 @@ def _enable_maintenance_events(
             not self.maintenance_events_config
             or not self.maintenance_events_config.enabled
         ):
+            self._maintenance_event_connection_handler = None
             return
 
         # Set up pool handler if available
@@ -702,7 +701,7 @@ def on_connect_check_health(self, check_health: bool = True):
                     "MAINT_NOTIFICATIONS",
                     "ON",
                     "moving-endpoint-type",
-                    endpoint_type,
+                    endpoint_type.value,
                     check_health=check_health,
                 )
                 response = self.read_response()

redis/maintenance_events.py

@@ -16,25 +16,18 @@ class MaintenanceState(enum.Enum):
     MAINTENANCE = "maintenance"
 
 
-class EndpointType:
-    """Constants for valid endpoint types used in CLIENT MAINT_NOTIFICATIONS command."""
+class EndpointType(enum.Enum):
+    """Valid endpoint types used in CLIENT MAINT_NOTIFICATIONS command."""
 
     INTERNAL_IP = "internal-ip"
     INTERNAL_FQDN = "internal-fqdn"
     EXTERNAL_IP = "external-ip"
     EXTERNAL_FQDN = "external-fqdn"
     NONE = "none"
 
-    @classmethod
-    def get_valid_types(cls):
-        """Return a set of all valid endpoint types."""
-        return {
-            cls.INTERNAL_IP,
-            cls.INTERNAL_FQDN,
-            cls.EXTERNAL_IP,
-            cls.EXTERNAL_FQDN,
-            cls.NONE,
-        }
+    def __str__(self):
+        """Return the string value of the enum."""
+        return self.value
 
 
 if TYPE_CHECKING:
@@ -438,7 +431,7 @@ def __init__(
         enabled: bool = True,
         proactive_reconnect: bool = True,
         relax_timeout: Optional[Number] = 20,
-        endpoint_type: Optional[str] = None,
+        endpoint_type: Optional[EndpointType] = None,
     ):
         """
         Initialize a new MaintenanceEventsConfig.
@@ -450,8 +443,7 @@ def __init__(
                 Defaults to True.
             relax_timeout (Number): The relax timeout to use for the connection during maintenance.
                 If -1 is provided - the relax timeout is disabled. Defaults to 20.
-            endpoint_type (Optional[str]): Override for the endpoint type to use in CLIENT MAINT_NOTIFICATIONS.
-                Must be one of: 'internal-ip', 'internal-fqdn', 'external-ip', 'external-fqdn', 'none'.
+            endpoint_type (Optional[EndpointType]): Override for the endpoint type to use in CLIENT MAINT_NOTIFICATIONS.
                 If None, the endpoint type will be automatically determined based on the host and TLS configuration.
                 Defaults to None.
 
@@ -461,19 +453,6 @@ def __init__(
         self.enabled = enabled
         self.relax_timeout = relax_timeout
         self.proactive_reconnect = proactive_reconnect
-
-        # Validate endpoint_type if provided
-        if (
-            endpoint_type is not None
-            and endpoint_type not in EndpointType.get_valid_types()
-        ):
-            valid_types = ", ".join(
-                f"'{t}'" for t in sorted(EndpointType.get_valid_types())
-            )
-            raise ValueError(
-                f"Invalid endpoint_type '{endpoint_type}'. Must be one of: {valid_types}"
-            )
-
         self.endpoint_type = endpoint_type
 
     def __repr__(self) -> str:
@@ -497,7 +476,9 @@ def is_relax_timeouts_enabled(self) -> bool:
         """
         return self.relax_timeout != -1
 
-    def get_endpoint_type(self, host: str, connection: "ConnectionInterface") -> str:
+    def get_endpoint_type(
+        self, host: str, connection: "ConnectionInterface"
+    ) -> EndpointType:
         """
         Determine the appropriate endpoint type for CLIENT MAINT_NOTIFICATIONS command.
 

tests/test_maintenance_events.py

@@ -674,23 +674,11 @@ class TestEndpointType:
 
     def test_endpoint_type_constants(self):
         """Test that the EndpointType constants are correct."""
-        assert EndpointType.INTERNAL_IP == "internal-ip"
-        assert EndpointType.INTERNAL_FQDN == "internal-fqdn"
-        assert EndpointType.EXTERNAL_IP == "external-ip"
-        assert EndpointType.EXTERNAL_FQDN == "external-fqdn"
-        assert EndpointType.NONE == "none"
-
-    def test_get_valid_types(self):
-        """Test that get_valid_types returns the expected set."""
-        valid_types = EndpointType.get_valid_types()
-        expected_types = {
-            "internal-ip",
-            "internal-fqdn",
-            "external-ip",
-            "external-fqdn",
-            "none",
-        }
-        assert valid_types == expected_types
+        assert EndpointType.INTERNAL_IP.value == "internal-ip"
+        assert EndpointType.INTERNAL_FQDN.value == "internal-fqdn"
+        assert EndpointType.EXTERNAL_IP.value == "external-ip"
+        assert EndpointType.EXTERNAL_FQDN.value == "external-fqdn"
+        assert EndpointType.NONE.value == "none"
 
 
 class TestMaintenanceEventsConfigEndpointType:
@@ -724,15 +712,10 @@ def get_resolved_ip(self):
 
     def test_config_validation_valid_endpoint_types(self):
         """Test that MaintenanceEventsConfig accepts valid endpoint types."""
-        for endpoint_type in EndpointType.get_valid_types():
+        for endpoint_type in EndpointType:
             config = MaintenanceEventsConfig(endpoint_type=endpoint_type)
             assert config.endpoint_type == endpoint_type
 
-    def test_config_validation_invalid_endpoint_type(self):
-        """Test that MaintenanceEventsConfig raises ValueError for invalid endpoint type."""
-        with pytest.raises(ValueError, match="Invalid endpoint_type"):
-            MaintenanceEventsConfig(endpoint_type="invalid-type")
-
     def test_config_validation_none_endpoint_type(self):
         """Test that MaintenanceEventsConfig accepts None as endpoint type."""
         config = MaintenanceEventsConfig(endpoint_type=None)