89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
from secrets import token_bytes
|
|
|
|
import pytest
|
|
|
|
from src.simulator.simulator_protocol import FarmNewBlockProtocol
|
|
from src.types.peer_info import PeerInfo
|
|
from src.util.ints import uint16, uint32, uint64
|
|
from tests.setup_nodes import setup_simulators_and_wallets
|
|
from src.consensus.block_rewards import calculate_base_fee, calculate_block_reward
|
|
from src.wallet.cc_wallet.cc_wallet import CCWallet
|
|
from tests.time_out_assert import time_out_assert
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def event_loop():
|
|
loop = asyncio.get_event_loop()
|
|
yield loop
|
|
|
|
|
|
class TestCCWalletBackup:
|
|
@pytest.fixture(scope="function")
|
|
async def two_wallet_nodes(self):
|
|
async for _ in setup_simulators_and_wallets(
|
|
1, 1, {"COINBASE_FREEZE_PERIOD": 0}
|
|
):
|
|
yield _
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_coin_backup(self, two_wallet_nodes):
|
|
num_blocks = 5
|
|
full_nodes, wallets = two_wallet_nodes
|
|
full_node_1, server_1 = full_nodes[0]
|
|
wallet_node, server_2 = wallets[0]
|
|
wallet = wallet_node.wallet_state_manager.main_wallet
|
|
|
|
ph = await wallet.get_new_puzzlehash()
|
|
|
|
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
|
|
for i in range(1, 4):
|
|
await full_node_1.farm_new_block(FarmNewBlockProtocol(ph))
|
|
|
|
funds = sum(
|
|
[
|
|
calculate_base_fee(uint32(i)) + calculate_block_reward(uint32(i))
|
|
for i in range(1, 4 - 2)
|
|
]
|
|
)
|
|
|
|
await time_out_assert(15, wallet.get_confirmed_balance, funds)
|
|
|
|
cc_wallet: CCWallet = await CCWallet.create_new_cc(
|
|
wallet_node.wallet_state_manager, wallet, uint64(100)
|
|
)
|
|
|
|
for i in range(1, num_blocks):
|
|
await full_node_1.farm_new_block(FarmNewBlockProtocol(ph))
|
|
|
|
await time_out_assert(15, cc_wallet.get_confirmed_balance, 100)
|
|
await time_out_assert(15, cc_wallet.get_unconfirmed_balance, 100)
|
|
|
|
# Write backup to file
|
|
filename = f"test-backup-{token_bytes(16).hex()}"
|
|
file_path = Path(filename)
|
|
await wallet_node.wallet_state_manager.create_wallet_backup(file_path)
|
|
|
|
# Close wallet and restart
|
|
db_path = wallet_node.wallet_state_manager.db_path
|
|
wallet_node._close()
|
|
await wallet_node._await_closed()
|
|
|
|
db_path.unlink()
|
|
|
|
started = await wallet_node._start()
|
|
assert started is False
|
|
|
|
await wallet_node._start(backup_file=file_path)
|
|
await server_2.start_client(PeerInfo("localhost", uint16(server_1._port)), None)
|
|
|
|
all_wallets = wallet_node.wallet_state_manager.wallets
|
|
assert len(all_wallets) == 2
|
|
|
|
cc_wallet_from_backup = wallet_node.wallet_state_manager.wallets[2]
|
|
|
|
await time_out_assert(15, cc_wallet_from_backup.get_confirmed_balance, 100)
|
|
if file_path.exists():
|
|
file_path.unlink()
|