From c7967ebeabde153d23c0c3fa5c1d5b9c63d2cccc Mon Sep 17 00:00:00 2001
From: opendansor <daniel@opentensor.dev>
Date: Fri, 31 May 2024 13:55:14 -0700
Subject: [PATCH] Fix return of ip version

---
 bittensor/axon.py                  |  3 +-
 tests/unit_tests/test_axon.py      | 50 ++++++++++++++++++++++++++++++
 tests/unit_tests/test_subtensor.py |  6 ++--
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/bittensor/axon.py b/bittensor/axon.py
index 476ed52db2..273cd79607 100644
--- a/bittensor/axon.py
+++ b/bittensor/axon.py
@@ -57,6 +57,7 @@
     SynapseException,
 )
 from bittensor.threadpool import PriorityThreadPoolExecutor
+from bittensor.utils import networking
 
 
 class FastAPIThreadedServer(uvicorn.Server):
@@ -393,7 +394,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
index ec0d9f5e53..36d0ba2d00 100644
--- 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
index 3709ec190d..c002aea66b 100644
--- 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