Merge pull request #1964 from opentensor/bugfix/opendansor/fix-ip-ver…

…sion-2 Fix return of ip version
opentensor · Jun 3, 2024 · 6899884 · 6899884
2 parents d88a743 + 923798b
commit 6899884
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 4 deletions.
diff --git a/bittensor/axon.py b/bittensor/axon.py
@@ -56,6 +56,7 @@
     SynapseException,
 )
 from bittensor.threadpool import PriorityThreadPoolExecutor
+from bittensor.utils import networking
 
 
 class FastAPIThreadedServer(uvicorn.Server):
@@ -392,7 +393,7 @@ def info(self) -> "bittensor.AxonInfo":
         return bittensor.AxonInfo(
             version=bittensor.__version_as_int__,
             ip=self.external_ip,
-            ip_type=4,
+            ip_type=networking.ip_version(self.external_ip),
             port=self.external_port,
             hotkey=self.wallet.hotkey.ss58_address,
             coldkey=self.wallet.coldkeypub.ss58_address,

diff --git a/tests/unit_tests/test_axon.py b/tests/unit_tests/test_axon.py
@@ -25,6 +25,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 # Third Party
+import netaddr
 import pytest
 from starlette.requests import Request
 from fastapi.testclient import TestClient
@@ -349,6 +350,55 @@ def test_to_string(info_return, expected_output, test_id):
         assert output == expected_output, f"Test ID: {test_id}"
 
 
+@pytest.mark.parametrize(
+    "ip, port, expected_ip_type, test_id",
+    [
+        # Happy path
+        (
+            "127.0.0.1",
+            8080,
+            4,
+            "valid_ipv4",
+        ),
+        (
+            "2001:0db8:85a3:0000:0000:8a2e:0370:7334",
+            3030,
+            6,
+            "valid_ipv6",
+        ),
+    ],
+)
+def test_valid_ipv4_and_ipv6_address(ip, port, expected_ip_type, test_id):
+    # Arrange
+    axon = Axon()
+    axon.ip = ip
+    axon.external_ip = ip
+    axon.port = port
+
+    # Act
+    ip_type = axon.info().ip_type
+
+    # Assert
+    assert ip_type == expected_ip_type, f"Test ID: {test_id}"
+
+
+@pytest.mark.parametrize(
+    "ip, port, expected_exception",
+    [
+        (
+            "This Is not a valid address",
+            65534,
+            netaddr.core.AddrFormatError,
+        ),
+    ],
+    ids=["failed to detect a valid IP " "address from %r"],
+)
+def test_invalid_ip_address(ip, port, expected_exception):
+    # Assert
+    with pytest.raises(expected_exception):
+        Axon(ip=ip, external_ip=ip, port=port).info()
+
+
 @pytest.mark.parametrize(
     "ip, port, ss58_address, started, forward_fns, expected_str, test_id",
     [

diff --git a/tests/unit_tests/test_subtensor.py b/tests/unit_tests/test_subtensor.py
@@ -34,8 +34,8 @@
 
 
 def test_serve_axon_with_external_ip_set():
-    internal_ip: str = "this is an internal ip"
-    external_ip: str = "this is an external ip"
+    internal_ip: str = "192.0.2.146"
+    external_ip: str = "2001:0db8:85a3:0000:0000:8a2e:0370:7334"
 
     mock_serve_axon = MagicMock(return_value=True)
 
@@ -76,7 +76,7 @@ def test_serve_axon_with_external_ip_set():
 
 
 def test_serve_axon_with_external_port_set():
-    external_ip: str = "this is an external ip"
+    external_ip: str = "2001:0db8:85a3:0000:0000:8a2e:0370:7334"
 
     internal_port: int = 1234
     external_port: int = 5678