async def get_stake_for_uid(self, uid) -> Balance:
stake = await self.substrate.get_runtime_state(
module='SubtensorModule', storage_function='Stake', params=[uid])
if not stake:
return Balance(0)
return Balance(stake['result'])
async def get_balance(self, address):
logger.debug("Getting balance for: {}", address)
result = await self.substrate.get_runtime_state(
module='System',
storage_function='Account',
params=[address],
block_hash=None)
balance_info = result.get('result')
if not balance_info:
logger.debug("{} has no balance", address)
return Balance(0)
balance = balance_info['data']['free']
logger.debug("{} has {} rao", address, balance)
return Balance(balance)
async def test_get_balance_no_balance(setup_chain):
hotkeypair = generate_keypair()
client = connect(hotkeypair,setup_chain)
await client.is_connected()
result = await client.get_balance(hotkeypair.ss58_address)
assert result == Balance(0)
async def add_stake(port, coldkeypair : Keypair, hotkeypair : Keypair, amount : Balance):
client = connect(coldkeypair, port)
await client.is_connected()
# Get the uid of the new neuron
uid = await client.get_uid_for_pubkey(hotkeypair.public_key)
assert uid is not None
# Get the amount of stake, should be 0
result = await client.get_stake_for_uid(uid)
assert int(result) == int(Balance(0))
# Add stake to new neuron
await client.add_stake(amount, hotkeypair.public_key)
async def test_add_stake_success(setup_chain):
coldkeypair = Keypair.create_from_uri("//Alice")
hotkeypair = generate_keypair()
client = connect(hotkeypair, setup_chain)
await client.is_connected()
# Subscribe a new neuron with the hotkey
await client.subscribe("8.8.8.8", 6666, 0, coldkeypair.public_key)
await asyncio.sleep(10)
# Now switch the connection the use the coldkey
client = connect(coldkeypair, setup_chain)
await client.is_connected()
# Get the uid of the new neuron
uid = await client.get_uid_for_pubkey(hotkeypair.public_key)
assert uid is not None
# Get the amount of stake, should be 0
result = await client.get_stake_for_uid(uid)
assert int (result) == int(Balance(0))
# Add stake to new neuron
result = await client.add_stake(Balance(4000), hotkeypair.public_key)
assert result is not None
assert 'extrinsic_hash' in result
# Wait for the extrinsic to complete
await asyncio.sleep(10)
# Get the amount of stake
result = await client.get_stake_for_uid(uid)
assert int(result) == int(Balance(4000))
async def test_get_stake_for_uid___has_stake(setup_chain):
hotkeyp = generate_keypair()
coldkeyp = Keypair.create_from_uri("//Alice")
client = connect(hotkeyp,setup_chain)
await client.is_connected()
await client.subscribe("8.8.8.8", 667, 0, coldkeyp.public_key)
await asyncio.sleep(10)
uid = await client.get_uid_for_pubkey(hotkeyp.public_key)
await add_stake(setup_chain, coldkeyp, hotkeyp, Balance(4000))
await asyncio.sleep(10)
result = await client.get_stake_for_uid(uid)
assert int(result) == 4000
async def test_transfer_success(setup_chain):
coldkey_alice = Keypair.create_from_uri("//Alice")
coldkey_bob = Keypair.create_from_uri("//Bob")
client = connect(coldkey_alice, setup_chain)
await client.is_connected()
balance_alice = await client.get_balance(coldkey_alice.ss58_address)
balance_bob = await client.get_balance(coldkey_bob.ss58_address)
result = await client.transfer(coldkey_bob.public_key, Balance(pow(10, 4)))
assert result is not None
assert 'extrinsic_hash' in result
# Wait until extrinsic is processed
await asyncio.sleep(10)
balance_alice_new = await client.get_balance(coldkey_alice.ss58_address)
balance_bob_new = await client.get_balance(coldkey_bob.ss58_address)
assert balance_alice_new < balance_alice
assert balance_bob_new > balance_bob
async def test_unstake_success(setup_chain):
coldkeypair = Keypair.create_from_uri("//Alice")
hotkeypair = generate_keypair()
hotkey_client = connect(hotkeypair,setup_chain)
await hotkey_client.is_connected()
# Subscribe a new neuron with the hotkey
await hotkey_client.subscribe("8.8.8.8", 6666, 0, coldkeypair.public_key)
await asyncio.sleep(10)
# Now switch the connection the use the coldkey
coldkey_client = connect(coldkeypair, setup_chain)
await coldkey_client.is_connected()
# Get the uid of the new neuron
uid = await coldkey_client.get_uid_for_pubkey(hotkeypair.public_key)
logger.error(uid)
assert uid is not None
# Get the amount of stake, should be 0
result = await coldkey_client.get_stake_for_uid(uid)
assert int(result) == int(Balance(0))
# Get the balance for the cold key, we use this for later comparison
balance = await coldkey_client.get_balance(coldkeypair.public_key)
# Add stake to new neuron
result = await coldkey_client.add_stake(Balance(4000), hotkeypair.public_key)
logger.info(result)
assert result is not None
assert 'extrinsic_hash' in result
# Wait for the extrinsic to complete
await asyncio.sleep(10)
# Get current balance, should be 4000 less than first balance
result = await coldkey_client.get_balance(coldkeypair.ss58_address)
assert int(result) == int(balance) - 4000
# Get the amount of stake, should be 4000
result = await coldkey_client.get_stake_for_uid(uid)
assert int(result) == int(Balance(4000))
# Now do the actual unstake
# Reconnect with coldkey account
coldkey_client = connect(coldkeypair, setup_chain)
await coldkey_client.is_connected()
# Do unstake
result = await coldkey_client.unstake(Balance(4000), hotkeypair.public_key)
assert result is not None
assert 'extrinsic_hash' in result
await asyncio.sleep(10)
# Check if balance is the same as what we started with
new_balance = await coldkey_client.get_balance(coldkeypair.ss58_address)
assert int(new_balance) == int(balance)