Applying review comments

Author: petyaslavova

redis/cluster.py

@@ -2199,31 +2199,17 @@ def _get_node_pubsub(self, node):
             return pubsub
 
     def _sharded_message_generator(self):
-        """
-        Iterate through pubsubs until a complete cycle is done.
-        """
-        while True:
+        for _ in range(len(self.node_pubsub_mapping)):
             pubsub = next(self._pubsubs_generator)
-
-            # None marks end of cycle
-            if pubsub is None:
-                break
-
             message = pubsub.get_message()
             if message is not None:
                 return message
-
         return None
 
     def _pubsubs_generator(self):
-        """
-        Generator that yields pubsubs in round-robin fashion.
-        Yields None to mark cycle boundaries.
-        """
         while True:
             current_nodes = list(self.node_pubsub_mapping.values())
             yield from current_nodes
-            yield None  # Cycle marker
 
     def get_sharded_message(
         self, ignore_subscribe_messages=False, timeout=0.0, target_node=None

tests/test_pubsub.py

@@ -959,20 +959,18 @@ def test_generator_handles_concurrent_mapping_changes(self, r):
         assert msg is not None
         assert msg["type"] == "ssubscribe"
 
-        # Get initial mapping size (if available)
-        initial_size = 0
-        if hasattr(pubsub, "node_pubsub_mapping"):
-            initial_size = len(pubsub.node_pubsub_mapping)
+        # Get initial mapping size (cluster pubsub only)
+        assert hasattr(pubsub, "node_pubsub_mapping"), "Test requires ClusterPubSub"
+        initial_size = len(pubsub.node_pubsub_mapping)
 
         # Subscribe to second channel (modifies mapping during potential iteration)
         pubsub.ssubscribe(channel2)
         msg = wait_for_message(pubsub, timeout=1.0, func=pubsub.get_sharded_message)
         assert msg is not None
         assert msg["type"] == "ssubscribe"
 
-        # Verify mapping was updated (if available)
-        if hasattr(pubsub, "node_pubsub_mapping"):
-            assert len(pubsub.node_pubsub_mapping) >= initial_size
+        # Verify mapping was updated
+        assert len(pubsub.node_pubsub_mapping) >= initial_size
 
         # Publish and read messages - should not raise RuntimeError
         r.spublish(channel1, "msg1")