Subscription manager bugfixes:

- Fix ``run_forever`` logic so that it doesn't require subscriptions to start
  and doesn't end when no subscriptions are present.

- Fix the ``unsubscribe`` method so that it can accept subscription ids.

- Fix the ``unsubscribe`` method so it behaves like the ``subscribe`` method,
  accepting both single or multiple subscriptions or subscription ids. Note
  this does change the signature from ``subscription`` to ``subscriptions``,
  but I think this can be a bugfix since the ``subscribe`` method accepts
  multiple subscriptions and ``unsubscribe`` should behave similarly.

- Fix ``subscribe`` so that it returns a list of subscription ids when a list
  of subscriptions is provided.

Add tests:

- Add a test for ``run_forever`` to make sure it can start without subscriptions
  and that the task does not stop when no subscriptions are present.

- Add a test for ``unsubscribe`` to make sure it can accept multiple subscriptions
  as hexstr ids or as subscription objects.

- Update subscriptions tests to check that a list of ids is returned when
  multiple subscriptions are provided to the ``subscribe`` method.

Author: fselmo

tests/core/subscriptions/test_subscription_manager.py

@@ -46,7 +46,8 @@ async def test_subscription_default_labels_are_unique(subscription_manager):
     sub3 = NewHeadsSubscription()
     sub4 = NewHeadsSubscription()
 
-    await subscription_manager.subscribe([sub1, sub2, sub3, sub4])
+    sub_ids = await subscription_manager.subscribe([sub1, sub2, sub3, sub4])
+    assert sub_ids == ["0x0", "0x1", "0x2", "0x3"]
 
     assert sub1.label != sub2.label != sub3.label != sub4.label
     assert sub1.label == "NewHeadsSubscription('newHeads',)"
@@ -133,3 +134,23 @@ async def test_unsubscribe_all_clears_all_subscriptions(subscription_manager):
     assert sub_container.subscriptions == []
     assert sub_container.subscriptions_by_id == {}
     assert sub_container.subscriptions_by_label == {}
+
+
+@pytest.mark.asyncio
+async def test_unsubscribe_with_hex_ids(subscription_manager):
+    sub1 = NewHeadsSubscription()
+    sub2 = PendingTxSubscription()
+    sub3 = NewHeadsSubscription()
+    sub4 = LogsSubscription()
+
+    sub_id1, sub_id2, sub_id3, sub_id4 = await subscription_manager.subscribe(
+        [sub1, sub2, sub3, sub4]
+    )
+
+    assert subscription_manager.subscriptions == [sub1, sub2, sub3, sub4]
+
+    assert await subscription_manager.unsubscribe(sub_id1) is True
+    assert subscription_manager.subscriptions == [sub2, sub3, sub4]
+
+    assert await subscription_manager.unsubscribe([sub_id2, sub_id3]) is True
+    assert subscription_manager.subscriptions == [sub4]

web3/_utils/module_testing/persistent_connection_provider.py

@@ -876,3 +876,41 @@ async def unsubscribe_subs(
 
         assert_no_subscriptions_left(sub_manager._subscription_container)
         await clean_up_task(unsubscribe_task)
+
+    @pytest.mark.asyncio
+    async def test_run_forever_starts_with_0_subs_and_runs_until_task_cancelled(
+        self, async_w3: AsyncWeb3
+    ) -> None:
+        sub_manager = async_w3.subscription_manager
+        assert_no_subscriptions_left(sub_manager._subscription_container)
+
+        run_forever_task = asyncio.create_task(
+            sub_manager.handle_subscriptions(run_forever=True)
+        )
+
+        await asyncio.sleep(0.1)
+        assert run_forever_task.done() is False
+        assert sub_manager.subscriptions == []
+
+        # subscribe to newHeads and validate it
+        new_heads_handler_test = SubscriptionHandlerTest()
+        sub1 = NewHeadsSubscription(
+            label="foo",
+            handler=new_heads_handler,
+            handler_context={"new_heads_handler_test": new_heads_handler_test},
+        )
+        sub_id = await sub_manager.subscribe(sub1)
+        assert is_hexstr(sub_id)
+        assert len(sub_manager.subscriptions) == 1
+        assert sub_manager.subscriptions[0] == sub1
+
+        # wait for the handler to unsubscribe
+        while sub_manager.subscriptions:
+            await asyncio.sleep(0.1)
+
+        assert new_heads_handler_test.passed
+        assert run_forever_task.done() is False
+        assert run_forever_task.cancelled() is False
+
+        # cleanup
+        await clean_up_task(run_forever_task)

web3/providers/persistent/subscription_manager.py

@@ -133,34 +133,88 @@ async def subscribe(
 
             sub_ids: List[HexStr] = []
             for sub in subscriptions:
-                await self.subscribe(sub)
+                sub_ids.append(await self.subscribe(sub))
             return sub_ids
         raise Web3TypeError("Expected a Subscription or a sequence of Subscriptions.")
 
-    async def unsubscribe(self, subscription: EthSubscription[Any]) -> bool:
+    @overload
+    async def unsubscribe(self, subscriptions: EthSubscription[Any]) -> bool:
+        ...
+
+    @overload
+    async def unsubscribe(self, subscriptions: HexStr) -> bool:
+        ...
+
+    @overload
+    async def unsubscribe(
+        self,
+        subscriptions: Sequence[Union[EthSubscription[Any], HexStr]],
+    ) -> bool:
+        ...
+
+    async def unsubscribe(
+        self,
+        subscriptions: Union[
+            EthSubscription[Any],
+            HexStr,
+            Sequence[Union[EthSubscription[Any], HexStr]],
+        ],
+    ) -> bool:
         """
-        Used to unsubscribe from a subscription.
+        Used to unsubscribe from one or multiple subscriptions.
 
-        :param subscription: The subscription to unsubscribe from.
-        :type subscription: EthSubscription
-        :return: ``True`` if unsubscribing was successful, ``False`` otherwise.
+        :param subscriptions: The subscription(s) to unsubscribe from.
+        :type subscriptions: Union[EthSubscription, Sequence[EthSubscription], HexStr,
+            Sequence[HexStr]]
+        :return: ``True`` if unsubscribing to all was successful, ``False`` otherwise
+            with a warning.
         :rtype: bool
         """
-        if subscription not in self.subscriptions:
-            raise Web3ValueError(
-                "Subscription not found or is not being managed by the subscription "
-                f"manager.\n    label: {subscription.label}\n    id: {subscription._id}"
-            )
-        if await self._w3.eth._unsubscribe(subscription.id):
-            self._remove_subscription(subscription)
-            self.logger.info(
-                "Successfully unsubscribed from subscription:\n    "
-                f"label: {subscription.label}\n    id: {subscription.id}"
-            )
-            if len(self._subscription_container.handler_subscriptions) == 0:
-                queue = self._provider._request_processor._handler_subscription_queue
-                await queue.put(SubscriptionProcessingFinished())
-            return True
+        if isinstance(subscriptions, EthSubscription) or isinstance(subscriptions, str):
+            if isinstance(subscriptions, str):
+                subscription_id = subscriptions
+                subscriptions = self.get_by_id(subscription_id)
+                if subscriptions is None:
+                    raise Web3ValueError(
+                        "Subscription not found or is not being managed by the "
+                        f"subscription manager.\n    id: {subscription_id}"
+                    )
+
+            if subscriptions not in self.subscriptions:
+                raise Web3ValueError(
+                    "Subscription not found or is not being managed by the "
+                    "subscription manager.\n    "
+                    f"label: {subscriptions.label}\n    id: {subscriptions._id}"
+                )
+
+            if await self._w3.eth._unsubscribe(subscriptions.id):
+                self._remove_subscription(subscriptions)
+                self.logger.info(
+                    "Successfully unsubscribed from subscription:\n    "
+                    f"label: {subscriptions.label}\n    id: {subscriptions.id}"
+                )
+
+                if len(self._subscription_container.handler_subscriptions) == 0:
+                    queue = (
+                        self._provider._request_processor._handler_subscription_queue
+                    )
+                    await queue.put(SubscriptionProcessingFinished())
+                return True
+
+        elif isinstance(subscriptions, Sequence):
+            if len(subscriptions) == 0:
+                raise Web3ValueError("No subscriptions provided.")
+
+            unsubscribed: List[bool] = []
+            for sub in subscriptions:
+                if isinstance(sub, str):
+                    sub = HexStr(sub)
+                unsubscribed.append(await self.unsubscribe(sub))
+            return all(unsubscribed)
+
+        self.logger.warning(
+            f"Failed to unsubscribe from subscription\n    subscription={subscriptions}"
+        )
         return False
 
     async def unsubscribe_all(self) -> bool:
@@ -195,15 +249,15 @@ async def handle_subscriptions(self, run_forever: bool = False) -> None:
         :type run_forever: bool
         :return: None
         """
-        if not self._subscription_container.handler_subscriptions:
+        if not self._subscription_container.handler_subscriptions and not run_forever:
             self.logger.warning(
                 "No handler subscriptions found. Subscription handler did not run."
             )
             return
 
         queue = self._provider._request_processor._handler_subscription_queue
-        try:
-            while run_forever or self._subscription_container.handler_subscriptions:
+        while run_forever or self._subscription_container.handler_subscriptions:
+            try:
                 response = cast(RPCResponse, await queue.get())
                 formatted_sub_response = cast(
                     FormattedEthSubscriptionResponse,
@@ -225,18 +279,20 @@ async def handle_subscriptions(self, run_forever: bool = False) -> None:
                             **sub._handler_context,
                         )
                     )
-        except SubscriptionProcessingFinished:
-            self.logger.info(
-                "All handler subscriptions have been unsubscribed from. "
-                "Stopping subscription handling."
-            )
-        except TaskNotRunning:
-            await asyncio.sleep(0)
-            self._provider._handle_listener_task_exceptions()
-            self.logger.error(
-                "Message listener background task for the provider has stopped "
-                "unexpectedly. Stopping subscription handling."
-            )
+            except SubscriptionProcessingFinished:
+                if not run_forever:
+                    self.logger.info(
+                        "All handler subscriptions have been unsubscribed from. "
+                        "Stopping subscription handling."
+                    )
+                    break
+            except TaskNotRunning:
+                await asyncio.sleep(0)
+                self._provider._handle_listener_task_exceptions()
+                self.logger.error(
+                    "Message listener background task for the provider has stopped "
+                    "unexpectedly. Stopping subscription handling."
+                )
 
         # no active handler subscriptions, clear the handler subscription queue
         self._provider._request_processor._reset_handler_subscription_queue()