250 lines
9.5 KiB
Python
250 lines
9.5 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
from secrets import token_bytes
|
|
import aiosqlite
|
|
import pytest
|
|
from src.util.ints import uint32, uint64, uint128
|
|
from src.wallet.wallet_store import WalletStore
|
|
from src.wallet.wallet_coin_record import WalletCoinRecord
|
|
from src.wallet.util.wallet_types import WalletType
|
|
from src.wallet.block_record import BlockRecord
|
|
from src.types.coin import Coin
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def event_loop():
|
|
loop = asyncio.get_event_loop()
|
|
yield loop
|
|
|
|
|
|
class TestWalletStore:
|
|
@pytest.mark.asyncio
|
|
async def test_store(self):
|
|
db_filename = Path("blockchain_wallet_store_test.db")
|
|
|
|
if db_filename.exists():
|
|
db_filename.unlink()
|
|
|
|
db_connection = await aiosqlite.connect(db_filename)
|
|
store = await WalletStore.create(db_connection)
|
|
try:
|
|
coin_1 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
coin_2 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
coin_3 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
coin_4 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
record_replaced = WalletCoinRecord(
|
|
coin_1, uint32(8), uint32(0), False, True, WalletType.STANDARD_WALLET, 0
|
|
)
|
|
record_1 = WalletCoinRecord(
|
|
coin_1, uint32(4), uint32(0), False, True, WalletType.STANDARD_WALLET, 0
|
|
)
|
|
record_2 = WalletCoinRecord(
|
|
coin_2, uint32(5), uint32(0), False, True, WalletType.STANDARD_WALLET, 0
|
|
)
|
|
record_3 = WalletCoinRecord(
|
|
coin_3,
|
|
uint32(5),
|
|
uint32(10),
|
|
True,
|
|
False,
|
|
WalletType.STANDARD_WALLET,
|
|
0,
|
|
)
|
|
record_4 = WalletCoinRecord(
|
|
coin_4,
|
|
uint32(5),
|
|
uint32(15),
|
|
True,
|
|
False,
|
|
WalletType.STANDARD_WALLET,
|
|
0,
|
|
)
|
|
|
|
# Test add (replace) and get
|
|
assert await store.get_coin_record(coin_1.name()) is None
|
|
await store.add_coin_record(record_replaced)
|
|
await store.add_coin_record(record_1)
|
|
await store.add_coin_record(record_2)
|
|
await store.add_coin_record(record_3)
|
|
await store.add_coin_record(record_4)
|
|
assert await store.get_coin_record(coin_1.name()) == record_1
|
|
|
|
# Test persistance
|
|
await db_connection.close()
|
|
db_connection = await aiosqlite.connect(db_filename)
|
|
store = await WalletStore.create(db_connection)
|
|
assert await store.get_coin_record(coin_1.name()) == record_1
|
|
|
|
# Test set spent
|
|
await store.set_spent(coin_1.name(), uint32(12))
|
|
assert (await store.get_coin_record(coin_1.name())).spent
|
|
assert (await store.get_coin_record(coin_1.name())).spent_block_index == 12
|
|
|
|
# No coins at height 3
|
|
assert len(await store.get_unspent_coins_at_height(3)) == 0
|
|
assert len(await store.get_unspent_coins_at_height(4)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(5)) == 4
|
|
assert len(await store.get_unspent_coins_at_height(11)) == 3
|
|
assert len(await store.get_unspent_coins_at_height(12)) == 2
|
|
assert len(await store.get_unspent_coins_at_height(15)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(16)) == 1
|
|
assert len(await store.get_unspent_coins_at_height()) == 1
|
|
|
|
assert len(await store.get_unspent_coins_for_wallet(0)) == 1
|
|
assert len(await store.get_unspent_coins_for_wallet(1)) == 0
|
|
|
|
coin_5 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
record_5 = WalletCoinRecord(
|
|
coin_5,
|
|
uint32(5),
|
|
uint32(15),
|
|
False,
|
|
False,
|
|
WalletType.STANDARD_WALLET,
|
|
1,
|
|
)
|
|
await store.add_coin_record(record_5)
|
|
assert len(await store.get_unspent_coins_for_wallet(1)) == 1
|
|
|
|
assert len(await store.get_spendable_for_index(100, 1)) == 1
|
|
assert len(await store.get_spendable_for_index(100, 0)) == 1
|
|
assert len(await store.get_spendable_for_index(0, 0)) == 0
|
|
|
|
coin_6 = Coin(token_bytes(32), coin_4.puzzle_hash, uint64(12312))
|
|
await store.add_coin_record(record_5)
|
|
record_6 = WalletCoinRecord(
|
|
coin_6,
|
|
uint32(5),
|
|
uint32(15),
|
|
True,
|
|
False,
|
|
WalletType.STANDARD_WALLET,
|
|
2,
|
|
)
|
|
await store.add_coin_record(record_6)
|
|
assert (
|
|
len(
|
|
await store.get_coin_records_by_puzzle_hash(
|
|
record_6.coin.puzzle_hash
|
|
)
|
|
)
|
|
== 2
|
|
) # 4 and 6
|
|
assert (
|
|
len(await store.get_coin_records_by_puzzle_hash(token_bytes(32))) == 0
|
|
)
|
|
|
|
assert await store.get_coin_record_by_coin_id(coin_6.name()) == record_6
|
|
assert await store.get_coin_record_by_coin_id(token_bytes(32)) is None
|
|
|
|
# BLOCKS
|
|
assert len(await store.get_lca_path()) == 0
|
|
|
|
# NOT lca block
|
|
br_1 = BlockRecord(
|
|
token_bytes(32),
|
|
token_bytes(32),
|
|
uint32(0),
|
|
uint128(100),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
uint64(0),
|
|
)
|
|
assert await store.get_block_record(br_1.header_hash) is None
|
|
await store.add_block_record(br_1, False)
|
|
assert len(await store.get_lca_path()) == 0
|
|
assert await store.get_block_record(br_1.header_hash) == br_1
|
|
|
|
# LCA genesis
|
|
await store.add_block_record(br_1, True)
|
|
assert await store.get_block_record(br_1.header_hash) == br_1
|
|
assert len(await store.get_lca_path()) == 1
|
|
assert (await store.get_lca_path())[br_1.header_hash] == br_1
|
|
|
|
br_2 = BlockRecord(
|
|
token_bytes(32),
|
|
token_bytes(32),
|
|
uint32(1),
|
|
uint128(100),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
uint64(0),
|
|
)
|
|
await store.add_block_record(br_2, False)
|
|
assert len(await store.get_lca_path()) == 1
|
|
await store.add_block_to_path(br_2.header_hash)
|
|
assert len(await store.get_lca_path()) == 2
|
|
assert (await store.get_lca_path())[br_2.header_hash] == br_2
|
|
|
|
br_3 = BlockRecord(
|
|
token_bytes(32),
|
|
token_bytes(32),
|
|
uint32(2),
|
|
uint128(100),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
uint64(0),
|
|
)
|
|
await store.add_block_record(br_3, True)
|
|
assert len(await store.get_lca_path()) == 3
|
|
await store.remove_blocks_from_path(1)
|
|
assert len(await store.get_lca_path()) == 2
|
|
|
|
await store.rollback_lca_to_block(0)
|
|
assert len(await store.get_unspent_coins_at_height()) == 0
|
|
|
|
coin_7 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
coin_8 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
coin_9 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
coin_10 = Coin(token_bytes(32), token_bytes(32), uint64(12312))
|
|
record_7 = WalletCoinRecord(
|
|
coin_7, uint32(0), uint32(1), True, False, WalletType.STANDARD_WALLET, 1
|
|
)
|
|
record_8 = WalletCoinRecord(
|
|
coin_8, uint32(1), uint32(2), True, False, WalletType.STANDARD_WALLET, 1
|
|
)
|
|
record_9 = WalletCoinRecord(
|
|
coin_9, uint32(2), uint32(3), True, False, WalletType.STANDARD_WALLET, 1
|
|
)
|
|
record_10 = WalletCoinRecord(
|
|
coin_10,
|
|
uint32(3),
|
|
uint32(4),
|
|
True,
|
|
False,
|
|
WalletType.STANDARD_WALLET,
|
|
1,
|
|
)
|
|
|
|
await store.add_coin_record(record_7)
|
|
await store.add_coin_record(record_8)
|
|
await store.add_coin_record(record_9)
|
|
await store.add_coin_record(record_10)
|
|
assert len(await store.get_unspent_coins_at_height(0)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(1)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(2)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(3)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(4)) == 0
|
|
|
|
await store.add_block_record(br_2, True)
|
|
await store.add_block_record(br_3, True)
|
|
|
|
await store.rollback_lca_to_block(1)
|
|
|
|
assert len(await store.get_unspent_coins_at_height(0)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(1)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(2)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(3)) == 1
|
|
assert len(await store.get_unspent_coins_at_height(4)) == 1
|
|
|
|
except AssertionError:
|
|
await db_connection.close()
|
|
raise
|
|
await db_connection.close()
|