diff --git a/bittensor_cli/src/bittensor/utils.py b/bittensor_cli/src/bittensor/utils.py index f9ade4507..f8e322f01 100644 --- a/bittensor_cli/src/bittensor/utils.py +++ b/bittensor_cli/src/bittensor/utils.py @@ -10,6 +10,7 @@ from functools import partial import re +import aiohttp from async_substrate_interface import AsyncExtrinsicReceipt from bittensor_wallet import Wallet, Keypair from bittensor_wallet.utils import SS58_FORMAT @@ -1512,3 +1513,30 @@ async def print_extrinsic_id( f":white_heavy_check_mark: Your extrinsic has been included as {ext_id}" ) return + + +async def check_img_mimetype(img_url: str) -> tuple[bool, str, str]: + """ + Checks to see if the given URL is an image, as defined by its mimetype. + + Args: + img_url: the URL to check + + Returns: + tuple: + bool: True if the URL has a MIME type indicating image (e.g. 'image/...'), False otherwise. + str: MIME type of the URL. + str: error message if the URL could not be retrieved + + """ + async with aiohttp.ClientSession() as session: + try: + async with session.get(img_url) as response: + if response.status != 200: + return False, "", "Could not fetch image" + elif "image/" not in response.content_type: + return False, response.content_type, "" + else: + return True, response.content_type, "" + except aiohttp.ClientError: + return False, "", "Could not fetch image" diff --git a/bittensor_cli/src/commands/subnets/subnets.py b/bittensor_cli/src/commands/subnets/subnets.py index 664657986..2d093b276 100644 --- a/bittensor_cli/src/commands/subnets/subnets.py +++ b/bittensor_cli/src/commands/subnets/subnets.py @@ -38,6 +38,7 @@ json_console, get_hotkey_pub_ss58, print_extrinsic_id, + check_img_mimetype, ) if TYPE_CHECKING: @@ -2257,7 +2258,28 @@ async def set_identity( ) -> tuple[bool, Optional[str]]: """Set identity information for a subnet""" - if not await subtensor.subnet_exists(netuid): + if prompt and (logo_url := subnet_identity.get("logo_url")): + sn_exists, img_validation = await asyncio.gather( + subtensor.subnet_exists(netuid), + check_img_mimetype(subnet_identity["logo_url"]), + ) + img_valid, content_type, err_msg = img_validation + if not img_valid: + confirmation_msg = f"Are you sure you want to use [blue]{logo_url}[/blue] as your image URL?" + if err_msg: + if not Confirm.ask(f"{err_msg}\n{confirmation_msg}"): + return False, None + else: + if not Confirm.ask( + f"The provided image's MIME type is {content_type}, which is not recognized as a valid" + f" image MIME type.\n{confirmation_msg}" + ): + return False, None + + else: + sn_exists = await subtensor.subnet_exists(netuid) + + if not sn_exists: err_console.print(f"Subnet {netuid} does not exist") return False, None diff --git a/pyproject.toml b/pyproject.toml index 47d49b92a..4a396266d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,14 +16,14 @@ requires-python = ">=3.9,<3.14" dependencies = [ "wheel", "async-substrate-interface>=1.5.2", - "aiohttp~=3.10.2", + "aiohttp~=3.13", "backoff~=2.2.1", "GitPython>=3.0.0", "netaddr~=1.3.0", "numpy>=2.0.1,<3.0.0", "Jinja2", "pycryptodome>=3.0.0,<4.0.0", - "PyYAML~=6.0.1", + "PyYAML~=6.0", "rich>=13.7,<15.0", "scalecodec==1.2.12", "typer>=0.16", diff --git a/tests/e2e_tests/test_set_identity.py b/tests/e2e_tests/test_set_identity.py new file mode 100644 index 000000000..9c008cdcd --- /dev/null +++ b/tests/e2e_tests/test_set_identity.py @@ -0,0 +1,201 @@ +import json +from unittest.mock import MagicMock, AsyncMock, patch + + +""" +Verify commands: +* btcli s create +* btcli s set-identity +* btcli s get-identity +""" + + +def test_set_id(local_chain, wallet_setup): + """ + Tests that the user is prompted to confirm that the incorrect text/html URL is + indeed the one they wish to set as their logo URL, and that when the MIME type is 'image/jpeg' + they are not given this prompt. + """ + wallet_path_alice = "//Alice" + netuid = 2 + + # Create wallet for Alice + keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup( + wallet_path_alice + ) + # Register a subnet with sudo as Alice + result = exec_command_alice( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--subnet-name", + "Test Subnet", + "--repo", + "https://github.com/username/repo", + "--contact", + "alice@opentensor.dev", + "--url", + "https://testsubnet.com", + "--discord", + "alice#1234", + "--description", + "A test subnet for e2e testing", + "--additional-info", + "Created by Alice", + "--logo-url", + "https://testsubnet.com/logo.png", + "--no-prompt", + "--json-output", + ], + ) + result_output = json.loads(result.stdout) + assert result_output["success"] is True + + mock_response = MagicMock() + mock_response.status = 200 + mock_response.content_type = "text/html" # bad MIME type + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.__aexit__ = AsyncMock(return_value=None) + + mock_session = MagicMock() + mock_session.get = MagicMock(return_value=mock_response) + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=None) + + with patch("aiohttp.ClientSession", return_value=mock_session): + set_identity = exec_command_alice( + "subnets", + "set-identity", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--hotkey", + wallet_alice.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--netuid", + str(netuid), + "--subnet-name", + sn_name := "Test Subnet", + "--github-repo", + sn_github := "https://github.com/username/repo", + "--subnet-contact", + sn_contact := "alice@opentensor.dev", + "--subnet-url", + sn_url := "https://testsubnet.com", + "--discord", + sn_discord := "alice#1234", + "--description", + sn_description := "A test subnet for e2e testing", + "--logo-url", + sn_logo_url := "https://testsubnet.com/logo.png", + "--additional-info", + sn_add_info := "Created by Alice", + "--prompt", + ], + inputs=["Y", "Y"], + ) + assert ( + f"Are you sure you want to use {sn_logo_url} as your image URL?" + in set_identity.stdout + ) + get_identity = exec_command_alice( + "subnets", + "get-identity", + extra_args=[ + "--chain", + "ws://127.0.0.1:9945", + "--netuid", + netuid, + "--json-output", + ], + ) + get_identity_output = json.loads(get_identity.stdout) + assert get_identity_output["subnet_name"] == sn_name + assert get_identity_output["github_repo"] == sn_github + assert get_identity_output["subnet_contact"] == sn_contact + assert get_identity_output["subnet_url"] == sn_url + assert get_identity_output["discord"] == sn_discord + assert get_identity_output["description"] == sn_description + assert get_identity_output["logo_url"] == sn_logo_url + assert get_identity_output["additional"] == sn_add_info + + mock_response = MagicMock() + mock_response.status = 200 + mock_response.content_type = "image/jpeg" # good MIME type + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.__aexit__ = AsyncMock(return_value=None) + + mock_session = MagicMock() + mock_session.get = MagicMock(return_value=mock_response) + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=None) + with patch("aiohttp.ClientSession", return_value=mock_session): + set_identity = exec_command_alice( + "subnets", + "set-identity", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--hotkey", + wallet_alice.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--netuid", + str(netuid), + "--subnet-name", + sn_name := "Test Subnet", + "--github-repo", + sn_github := "https://github.com/username/repo", + "--subnet-contact", + sn_contact := "alice@opentensor.dev", + "--subnet-url", + sn_url := "https://testsubnet.com", + "--discord", + sn_discord := "alice#1234", + "--description", + sn_description := "A test subnet for e2e testing", + "--logo-url", + sn_logo_url := "https://testsubnet.com/logo.png", + "--additional-info", + sn_add_info := "Created by Alice", + "--prompt", + ], + inputs=["Y"], + ) + assert ( + f"Are you sure you want to use {sn_logo_url} as your image URL?" + not in set_identity.stdout + ) + get_identity = exec_command_alice( + "subnets", + "get-identity", + extra_args=[ + "--chain", + "ws://127.0.0.1:9945", + "--netuid", + netuid, + "--json-output", + ], + ) + get_identity_output = json.loads(get_identity.stdout) + assert get_identity_output["subnet_name"] == sn_name + assert get_identity_output["github_repo"] == sn_github + assert get_identity_output["subnet_contact"] == sn_contact + assert get_identity_output["subnet_url"] == sn_url + assert get_identity_output["discord"] == sn_discord + assert get_identity_output["description"] == sn_description + assert get_identity_output["logo_url"] == sn_logo_url + assert get_identity_output["additional"] == sn_add_info diff --git a/tests/e2e_tests/utils.py b/tests/e2e_tests/utils.py index 9917c2c10..323797356 100644 --- a/tests/e2e_tests/utils.py +++ b/tests/e2e_tests/utils.py @@ -28,8 +28,8 @@ def __call__( self, command: str, sub_command: str, - extra_args: Optional[list[str]], - inputs: Optional[list[str]], + extra_args: Optional[list[str]] = None, + inputs: Optional[list[str]] = None, ) -> Result: ... diff --git a/tests/unit_tests/test_utils.py b/tests/unit_tests/test_utils.py index 9aa737032..2f43b4605 100644 --- a/tests/unit_tests/test_utils.py +++ b/tests/unit_tests/test_utils.py @@ -1,5 +1,8 @@ from bittensor_cli.src.bittensor import utils import pytest +from unittest.mock import AsyncMock, patch, MagicMock + +from bittensor_cli.src.bittensor.utils import check_img_mimetype @pytest.mark.parametrize( @@ -27,3 +30,45 @@ ) def test_decode_hex_identity_dict(input_dict, expected_result): assert utils.decode_hex_identity_dict(input_dict) == expected_result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "img_url,status,content_type,expected_result", + [ + ( + "https://github.com/dougsillars/dougsillars/blob/main/twitter.jpg", + 200, + "text/html", + (False, "text/html", ""), + ), + ( + "https://raw.githubusercontent.com/dougsillars/dougsillars/refs/heads/main/twitter.jpg", + 200, + "image/jpeg", + (True, "image/jpeg", ""), + ), + ( + "https://abs-0.twimg.com/emoji/v2/svg/1f5fv.svg", + 404, + "", + (False, "", "Could not fetch image"), + ), + ], +) +async def test_get_image_url(img_url, status, content_type, expected_result): + mock_response = MagicMock() + mock_response.status = status + mock_response.content_type = content_type + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.__aexit__ = AsyncMock(return_value=None) + + # Create mock session + mock_session = MagicMock() + mock_session.get = MagicMock(return_value=mock_response) + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=None) + + # Patch ClientSession + with patch("aiohttp.ClientSession", return_value=mock_session): + assert await check_img_mimetype(img_url) == expected_result