chia-blockchain/tests/test_wallet.py

100 lines
3.2 KiB
Python

import asyncio
import pytest
from blspy import ExtendedPrivateKey
from src.protocols.wallet_protocol import RespondBody
from src.wallet.wallet import Wallet
from src.wallet.wallet_node import WalletNode
from tests.setup_nodes import setup_two_nodes, test_constants, bt
@pytest.fixture(scope="module")
def event_loop():
loop = asyncio.get_event_loop()
yield loop
class TestWallet:
@pytest.fixture(scope="function")
async def two_nodes(self):
async for _ in setup_two_nodes({"COINBASE_FREEZE_PERIOD": 0}):
yield _
@pytest.mark.asyncio
async def test_wallet_receive_body(self, two_nodes):
sk = bytes(ExtendedPrivateKey.from_seed(b"")).hex()
key_config = {"wallet_sk": sk}
wallet_node = await WalletNode.create({}, key_config)
wallet = wallet_node.wallet
await wallet_node.wallet_store._clear_database()
await wallet_node.tx_store._clear_database()
num_blocks = 10
blocks = bt.get_consecutive_blocks(
test_constants,
num_blocks,
[],
10,
reward_puzzlehash=wallet.get_new_puzzlehash(),
)
for i in range(1, num_blocks):
a = RespondBody(blocks[i].body, blocks[i].height)
await wallet_node.received_body(a)
assert await wallet.get_confirmed_balance() == 144000000000000
await wallet_node.wallet_store.close()
await wallet_node.tx_store.close()
@pytest.mark.asyncio
async def test_wallet_make_transaction(self, two_nodes):
sk = bytes(ExtendedPrivateKey.from_seed(b"")).hex()
sk_b = bytes(ExtendedPrivateKey.from_seed(b"b")).hex()
key_config = {"wallet_sk": sk}
key_config_b = {"wallet_sk": sk_b}
wallet_node = await WalletNode.create({}, key_config)
wallet = wallet_node.wallet
await wallet_node.wallet_store._clear_database()
await wallet_node.tx_store._clear_database()
wallet_node_b = await WalletNode.create({}, key_config_b)
wallet_b = wallet_node_b.wallet
await wallet_node_b.wallet_store._clear_database()
await wallet_node_b.tx_store._clear_database()
num_blocks = 10
blocks = bt.get_consecutive_blocks(
test_constants,
num_blocks,
[],
10,
reward_puzzlehash=wallet.get_new_puzzlehash(),
)
for i in range(1, num_blocks):
a = RespondBody(blocks[i].body, blocks[i].height)
await wallet_node.received_body(a)
assert await wallet.get_confirmed_balance() == 144000000000000
spend_bundle = await wallet.generate_signed_transaction(
10, wallet_b.get_new_puzzlehash(), 0
)
await wallet.push_transaction(spend_bundle)
confirmed_balance = await wallet.get_confirmed_balance()
unconfirmed_balance = await wallet.get_unconfirmed_balance()
assert confirmed_balance == 144000000000000
assert unconfirmed_balance == confirmed_balance - 10
await wallet_node.wallet_store.close()
await wallet_node.tx_store.close()
await wallet_node_b.wallet_store.close()
await wallet_node_b.tx_store.close()