Source code for bittensor.mock.wallet_mock

# The MIT License (MIT)

# Copyright © 2021 Yuma Rao
# Copyright © 2022 Opentensor Foundation
# Copyright © 2023 Opentensor Technologies

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
# documentation files (the “Software”), to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of
# the Software.

# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
# THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import os
import bittensor
from typing import Optional
from Crypto.Hash import keccak

from .keyfile_mock import MockKeyfile


[docs] class MockWallet(bittensor.wallet): """ Mocked Version of the bittensor wallet class, meant to be used for testing """ def __init__(self, **kwargs): r"""Init bittensor wallet object containing a hot and coldkey. Args: _mock (required=True, default=False): If true creates a mock wallet with random keys. """ super().__init__(**kwargs) # For mocking. self._is_mock = True self._mocked_coldkey_keyfile = None self._mocked_hotkey_keyfile = None @property def hotkey_file(self) -> "bittensor.keyfile": if self._is_mock: if self._mocked_hotkey_keyfile == None: self._mocked_hotkey_keyfile = MockKeyfile(path="MockedHotkey") return self._mocked_hotkey_keyfile else: wallet_path = os.path.expanduser(os.path.join(self.path, self.name)) hotkey_path = os.path.join(wallet_path, "hotkeys", self.hotkey_str) return bittensor.keyfile(path=hotkey_path) @property def coldkey_file(self) -> "bittensor.keyfile": if self._is_mock: if self._mocked_coldkey_keyfile == None: self._mocked_coldkey_keyfile = MockKeyfile(path="MockedColdkey") return self._mocked_coldkey_keyfile else: wallet_path = os.path.expanduser(os.path.join(self.path, self.name)) coldkey_path = os.path.join(wallet_path, "coldkey") return bittensor.keyfile(path=coldkey_path) @property def coldkeypub_file(self) -> "bittensor.keyfile": if self._is_mock: if self._mocked_coldkey_keyfile == None: self._mocked_coldkey_keyfile = MockKeyfile(path="MockedColdkeyPub") return self._mocked_coldkey_keyfile else: wallet_path = os.path.expanduser(os.path.join(self.path, self.name)) coldkeypub_path = os.path.join(wallet_path, "coldkeypub.txt") return bittensor.keyfile(path=coldkeypub_path)
[docs] def get_mock_wallet( coldkey: "bittensor.Keypair" = None, hotkey: "bittensor.Keypair" = None ): wallet = MockWallet(name="mock_wallet", hotkey="mock", path="/tmp/mock_wallet") if not coldkey: coldkey = bittensor.Keypair.create_from_mnemonic( bittensor.Keypair.generate_mnemonic() ) if not hotkey: hotkey = bittensor.Keypair.create_from_mnemonic( bittensor.Keypair.generate_mnemonic() ) wallet.set_coldkey(coldkey, encrypt=False, overwrite=True) wallet.set_coldkeypub(coldkey, encrypt=False, overwrite=True) wallet.set_hotkey(hotkey, encrypt=False, overwrite=True) return wallet
[docs] def get_mock_keypair(uid: int, test_name: Optional[str] = None) -> bittensor.Keypair: """ Returns a mock keypair from a uid and optional test_name. If test_name is not provided, the uid is the only seed. If test_name is provided, the uid is hashed with the test_name to create a unique seed for the test. """ if test_name is not None: hashed_test_name: bytes = keccak.new( digest_bits=256, data=test_name.encode("utf-8") ).digest() hashed_test_name_as_int: int = int.from_bytes( hashed_test_name, byteorder="big", signed=False ) uid = uid + hashed_test_name_as_int return bittensor.Keypair.create_from_seed( seed_hex=int.to_bytes(uid, 32, "big", signed=False), ss58_format=bittensor.__ss58_format__, )
[docs] def get_mock_hotkey(uid: int) -> str: return get_mock_keypair(uid).ss58_address
[docs] def get_mock_coldkey(uid: int) -> str: return get_mock_keypair(uid).ss58_address