chia-blockchain/tests/wallet/test_wallet_sync.py
Mariano Sorgente f772b4d22f Test passing
2020-09-16 01:19:28 -07:00

295 lines
12 KiB
Python

import asyncio
import pytest
from src.types.peer_info import PeerInfo
from src.protocols import full_node_protocol
from src.util.ints import uint16, uint64, uint32
from tests.setup_nodes import setup_node_and_wallet, test_constants, bt
from src.types.spend_bundle import SpendBundle
from src.full_node.bundle_tools import best_solution_program
from src.util.wallet_tools import WalletTool
from src.types.coin import Coin
from src.consensus.coinbase import create_coinbase_coin
from tests.time_out_assert import time_out_assert
def wallet_height_at_least(wallet_node, h):
if (
wallet_node.wallet_state_manager.block_records[
wallet_node.wallet_state_manager.lca
].height
>= h
):
return True
return False
@pytest.fixture(scope="module")
def event_loop():
loop = asyncio.get_event_loop()
yield loop
class TestWalletSync:
@pytest.fixture(scope="function")
async def wallet_node(self):
constants = test_constants.replace(COINBASE_FREEZE_PERIOD=0)
async for _ in setup_node_and_wallet(constants):
yield _
@pytest.fixture(scope="function")
async def wallet_node_starting_height(self):
async for _ in setup_node_and_wallet(test_constants, starting_height=100):
yield _
@pytest.mark.asyncio
async def test_basic_sync_wallet(self, wallet_node):
num_blocks = 300 # This must be greater than the short_sync in wallet_node
blocks = bt.get_consecutive_blocks(test_constants, num_blocks, [])
full_node_1, wallet_node, server_1, server_2 = wallet_node
for i in range(1, len(blocks)):
async for _ in full_node_1.respond_block(
full_node_protocol.RespondBlock(blocks[i])
):
pass
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
# The second node should eventually catch up to the first one, and have the
# same tip at height num_blocks - 1.
await time_out_assert(
200, wallet_height_at_least, True, wallet_node, num_blocks - 6
)
# Tests a reorg with the wallet
blocks_reorg = bt.get_consecutive_blocks(test_constants, 15, blocks[:-5])
for i in range(1, len(blocks_reorg)):
async for msg in full_node_1.respond_block(
full_node_protocol.RespondBlock(blocks_reorg[i])
):
server_1.push_message(msg)
await time_out_assert(200, wallet_height_at_least, True, wallet_node, 33)
@pytest.mark.asyncio
async def test_fast_sync_wallet(self, wallet_node_starting_height):
num_blocks = 25 # This must be greater than the short_sync in wallet_node
blocks = bt.get_consecutive_blocks(test_constants, num_blocks, [])
full_node_1, wallet_node, server_1, server_2 = wallet_node_starting_height
for i in range(1, len(blocks)):
async for _ in full_node_1.respond_block(
full_node_protocol.RespondBlock(blocks[i])
):
pass
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
await time_out_assert(
60, wallet_height_at_least, True, wallet_node, num_blocks - 6
)
@pytest.mark.asyncio
async def test_short_sync_wallet(self, wallet_node):
num_blocks = 5 # This must be lower than the short_sync in wallet_node
blocks = bt.get_consecutive_blocks(test_constants, num_blocks, [], 10)
full_node_1, wallet_node, server_1, server_2 = wallet_node
for i in range(1, len(blocks)):
async for _ in full_node_1.respond_block(
full_node_protocol.RespondBlock(blocks[i])
):
pass
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
await time_out_assert(60, wallet_height_at_least, True, wallet_node, 3)
@pytest.mark.asyncio
async def test_short_sync_with_transactions_wallet(self, wallet_node):
full_node_1, wallet_node, server_1, server_2 = wallet_node
wallet_a = wallet_node.wallet_state_manager.main_wallet
wallet_a_dummy = WalletTool()
wallet_b = WalletTool()
coinbase_puzzlehash = await wallet_a.get_new_puzzlehash()
coinbase_puzzlehash_rest = wallet_b.get_new_puzzlehash()
puzzle_hashes = [await wallet_a.get_new_puzzlehash() for _ in range(10)]
puzzle_hashes.append(wallet_b.get_new_puzzlehash())
blocks = bt.get_consecutive_blocks(
test_constants, 3, [], 10, b"", coinbase_puzzlehash
)
for block in blocks:
[
_
async for _ in full_node_1.respond_block(
full_node_protocol.RespondBlock(block)
)
]
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
await time_out_assert(60, wallet_height_at_least, True, wallet_node, 1)
server_2.global_connections.close_all_connections()
dic_h = {}
prev_coin = blocks[1].get_coinbase()
for i in range(11):
pk, sk = await wallet_a.wallet_state_manager.get_keys(prev_coin.puzzle_hash)
transaction_unsigned = wallet_a_dummy.generate_unsigned_transaction(
1000, puzzle_hashes[i], prev_coin, {}, 0, secretkey=sk
)
spend_bundle = await wallet_a.sign_transaction(transaction_unsigned)
block_spendbundle = SpendBundle.aggregate([spend_bundle])
program = best_solution_program(block_spendbundle)
aggsig = block_spendbundle.aggregated_signature
prev_coin = Coin(prev_coin.name(), puzzle_hashes[i], uint64(1000))
dic_h[i + 4] = (program, aggsig)
blocks = bt.get_consecutive_blocks(
test_constants, 13, blocks, 10, b"", coinbase_puzzlehash_rest, dic_h
)
# Move chain to height 16, with consecutive transactions in blocks 4 to 14
for block in blocks:
async for _ in full_node_1.respond_block(
full_node_protocol.RespondBlock(block)
):
pass
# Do a short sync from 0 to 14
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
await time_out_assert(60, wallet_height_at_least, True, wallet_node, 14)
server_2.global_connections.close_all_connections()
# 2 block rewards and 3 fees
assert await wallet_a.get_confirmed_balance() == (
blocks[1].get_coinbase().amount * 2
) + (blocks[1].get_fees_coin().amount * 3)
# All of our coins are spent and puzzle hashes present
for puzzle_hash in puzzle_hashes[:-1]:
records = await wallet_node.wallet_state_manager.wallet_store.get_coin_records_by_puzzle_hash(
puzzle_hash
)
assert len(records) == 1
assert records[0].spent and not records[0].coinbase
# Then do the same but in a reorg chain
dic_h = {}
prev_coin = blocks[1].get_coinbase()
for i in range(11):
pk, sk = await wallet_a.wallet_state_manager.get_keys(prev_coin.puzzle_hash)
transaction_unsigned = wallet_a_dummy.generate_unsigned_transaction(
1000, puzzle_hashes[i], prev_coin, {}, 0, secretkey=sk
)
spend_bundle = await wallet_a.sign_transaction(transaction_unsigned)
block_spendbundle = SpendBundle.aggregate([spend_bundle])
program = best_solution_program(block_spendbundle)
aggsig = block_spendbundle.aggregated_signature
prev_coin = Coin(prev_coin.name(), puzzle_hashes[i], uint64(1000))
dic_h[i + 4] = (program, aggsig)
blocks = bt.get_consecutive_blocks(
test_constants,
31,
blocks[:4],
10,
b"this is a reorg",
coinbase_puzzlehash_rest,
dic_h,
)
# Move chain to height 34, with consecutive transactions in blocks 4 to 14
for block in blocks:
async for _ in full_node_1.respond_block(
full_node_protocol.RespondBlock(block)
):
pass
# Do a sync from 0 to 22
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
await time_out_assert(60, wallet_height_at_least, True, wallet_node, 28)
server_2.global_connections.close_all_connections()
# 2 block rewards and 3 fees
assert await wallet_a.get_confirmed_balance() == (
blocks[1].get_coinbase().amount * 2
) + (blocks[1].get_fees_coin().amount * 3)
# All of our coins are spent and puzzle hashes present
for puzzle_hash in puzzle_hashes[:-1]:
records = await wallet_node.wallet_state_manager.wallet_store.get_coin_records_by_puzzle_hash(
puzzle_hash
)
assert len(records) == 1
assert records[0].spent and not records[0].coinbase
# Test spending the rewards earned in reorg
new_coinbase_puzzlehash = await wallet_a.get_new_puzzlehash()
another_puzzlehash = await wallet_a.get_new_puzzlehash()
dic_h = {}
pk, sk = await wallet_a.wallet_state_manager.get_keys(new_coinbase_puzzlehash)
coinbase_coin = create_coinbase_coin(
uint32(25), new_coinbase_puzzlehash, uint64(14000000000000)
)
transaction_unsigned = wallet_a_dummy.generate_unsigned_transaction(
7000000000000, another_puzzlehash, coinbase_coin, {}, 0, secretkey=sk
)
spend_bundle = await wallet_a.sign_transaction(transaction_unsigned)
block_spendbundle = SpendBundle.aggregate([spend_bundle])
program = best_solution_program(block_spendbundle)
aggsig = block_spendbundle.aggregated_signature
dic_h[26] = (program, aggsig)
# Farm a block (25) to ourselves
blocks = bt.get_consecutive_blocks(
test_constants,
1,
blocks[:25],
10,
b"this is yet another reorg",
new_coinbase_puzzlehash,
)
# Brings height up to 40, with block 31 having half our reward spent to us
blocks = bt.get_consecutive_blocks(
test_constants,
15,
blocks,
10,
b"this is yet another reorg more blocks",
coinbase_puzzlehash_rest,
dic_h,
)
for block in blocks:
async for _ in full_node_1.respond_block(
full_node_protocol.RespondBlock(block)
):
pass
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
await time_out_assert(60, wallet_height_at_least, True, wallet_node, 38)
# 2 block rewards and 4 fees, plus 7000000000000 coins
assert (
await wallet_a.get_confirmed_balance()
== (blocks[1].get_coinbase().amount * 2)
+ (blocks[1].get_fees_coin().amount * 4)
+ 7000000000000
)
records = await wallet_node.wallet_state_manager.wallet_store.get_coin_records_by_puzzle_hash(
new_coinbase_puzzlehash
)
# Fee and coinbase
assert len(records) == 2
print(records)
assert records[0].spent != records[1].spent
assert records[0].coinbase == records[1].coinbase
records = await wallet_node.wallet_state_manager.wallet_store.get_coin_records_by_puzzle_hash(
another_puzzlehash
)
assert len(records) == 1
assert not records[0].spent
assert not records[0].coinbase