From f33e45a40e2e06b6b7e88af50c93b8b6cbf86b08 Mon Sep 17 00:00:00 2001
From: Roman <roman@opentensor.dev>
Date: Fri, 20 Dec 2024 11:08:44 -0800
Subject: [PATCH 1/2] add connection limit error handler

---
 bittensor/core/subtensor.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py
index 1709701203..b0aa0e15c1 100644
--- a/bittensor/core/subtensor.py
+++ b/bittensor/core/subtensor.py
@@ -17,6 +17,7 @@
 from scalecodec.type_registry import load_type_registry_preset
 from scalecodec.types import ScaleType
 from substrateinterface.base import QueryMapResult, SubstrateInterface
+from websockets.exceptions import InvalidStatus
 from websockets.sync import client as ws_client
 
 from bittensor.core import settings
@@ -233,6 +234,7 @@ def _get_substrate(self, force: bool = False):
                     open_timeout=self._connection_timeout,
                     max_size=2**32,
                 )
+
             self.substrate = SubstrateInterface(
                 ss58_format=settings.SS58_FORMAT,
                 use_remote_preset=True,
@@ -244,19 +246,26 @@ def _get_substrate(self, force: bool = False):
                     f"Connected to {self.network} network and {self.chain_endpoint}."
                 )
 
-        except (ConnectionRefusedError, ssl.SSLError) as error:
-            logging.error(
-                f"<red>Could not connect to</red> <blue>{self.network}</blue> <red>network with</red> <blue>{self.chain_endpoint}</blue> <red>chain endpoint.</red>",
+        except ConnectionRefusedError as error:
+            logging.critical(
+                f"[red]Could not connect to[/red] [blue]{self.network}[/blue] [red]network with[/red] [blue]{self.chain_endpoint}[/blue] [red]chain endpoint.[/red]",
             )
             raise ConnectionRefusedError(error.args)
-        except ssl.SSLError as e:
+
+        except ssl.SSLError as error:
             logging.critical(
                 "SSL error occurred. To resolve this issue, run the following command in your terminal:"
             )
             logging.critical("[blue]sudo python -m bittensor certifi[/blue]")
             raise RuntimeError(
                 "SSL configuration issue, please follow the instructions above."
-            ) from e
+            ) from error
+
+        except InvalidStatus as error:
+            logging.critical(
+                f"[red]You have reached the limit of simultaneous connections to the server.[/red]"
+            )
+            raise InvalidStatus(error.response) from error
 
     @staticmethod
     def config() -> "Config":

From bd368dfff3dfeb501e8db2478b570d221f3e837e Mon Sep 17 00:00:00 2001
From: Roman <roman@opentensor.dev>
Date: Fri, 20 Dec 2024 11:09:06 -0800
Subject: [PATCH 2/2] add test

---
 tests/unit_tests/test_subtensor.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tests/unit_tests/test_subtensor.py b/tests/unit_tests/test_subtensor.py
index 5b25b62125..5b9e15339d 100644
--- a/tests/unit_tests/test_subtensor.py
+++ b/tests/unit_tests/test_subtensor.py
@@ -2880,3 +2880,27 @@ def test_set_weights_with_commit_reveal_enabled(subtensor, mocker):
         wait_for_finalization=fake_wait_for_finalization,
     )
     assert result == mocked_commit_reveal_v3_extrinsic.return_value
+
+
+def test_connection_limit(mocker):
+    """Test connection limit is not exceeded."""
+    # Technically speaking, this test should exist in integration tests. But to reduce server costs we will leave this
+    # test here.
+
+    # Preps
+    mocker.patch.object(
+        subtensor_module.ws_client,
+        "connect",
+        side_effect=subtensor_module.InvalidStatus(
+            response=mocker.Mock(
+                response=mocker.Mock(
+                    status_code=429, message="test connection limit error"
+                )
+            )
+        ),
+    )
+    # Call with assertions
+
+    with pytest.raises(subtensor_module.InvalidStatus):
+        for i in range(2):
+            Subtensor("test")