forked from Chia-Network/chia-blockchain
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_puzzle_store.py
104 lines (87 loc) · 3.97 KB
/
test_puzzle_store.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import asyncio
from pathlib import Path
from secrets import token_bytes
import aiosqlite
import pytest
from blspy import AugSchemeMPL
from chia.util.ints import uint32
from chia.wallet.derivation_record import DerivationRecord
from chia.wallet.util.wallet_types import WalletType
from chia.wallet.wallet_puzzle_store import WalletPuzzleStore
@pytest.fixture(scope="module")
def event_loop():
loop = asyncio.get_event_loop()
yield loop
class TestPuzzleStore:
@pytest.mark.asyncio
async def test_puzzle_store(self):
db_filename = Path("puzzle_store_test.db")
if db_filename.exists():
db_filename.unlink()
con = await aiosqlite.connect(db_filename)
db = await WalletPuzzleStore.create(con)
try:
derivation_recs = []
# wallet_types = [t for t in WalletType]
[t for t in WalletType]
for i in range(1000):
derivation_recs.append(
DerivationRecord(
uint32(i),
token_bytes(32),
AugSchemeMPL.key_gen(token_bytes(32)).get_g1(),
WalletType.STANDARD_WALLET,
uint32(1),
)
)
derivation_recs.append(
DerivationRecord(
uint32(i),
token_bytes(32),
AugSchemeMPL.key_gen(token_bytes(32)).get_g1(),
WalletType.RATE_LIMITED,
uint32(2),
)
)
assert await db.puzzle_hash_exists(derivation_recs[0].puzzle_hash) is False
assert await db.index_for_pubkey(derivation_recs[0].pubkey) is None
assert await db.index_for_puzzle_hash(derivation_recs[2].puzzle_hash) is None
assert await db.wallet_info_for_puzzle_hash(derivation_recs[2].puzzle_hash) is None
assert len((await db.get_all_puzzle_hashes())) == 0
assert await db.get_last_derivation_path() is None
assert await db.get_unused_derivation_path() is None
assert await db.get_derivation_record(0, 2) is None
await db.add_derivation_paths(derivation_recs)
assert await db.puzzle_hash_exists(derivation_recs[0].puzzle_hash) is True
phs_1 = [derivation_recs[0].puzzle_hash]
phs_2 = [32 * bytes([1]), derivation_recs[0].puzzle_hash]
phs_3 = [derivation_recs[0].puzzle_hash, 32 * bytes([1])]
phs_4 = [32 * bytes([1]), 32 * bytes([2])]
phs_5 = []
assert await db.one_of_puzzle_hashes_exists(phs_1) is True
assert await db.one_of_puzzle_hashes_exists(phs_2) is True
assert await db.one_of_puzzle_hashes_exists(phs_3) is True
assert await db.one_of_puzzle_hashes_exists(phs_4) is False
assert await db.one_of_puzzle_hashes_exists(phs_5) is False
assert await db.index_for_pubkey(derivation_recs[4].pubkey) == 2
assert await db.index_for_puzzle_hash(derivation_recs[2].puzzle_hash) == 1
assert await db.wallet_info_for_puzzle_hash(derivation_recs[2].puzzle_hash) == (
derivation_recs[2].wallet_id,
derivation_recs[2].wallet_type,
)
assert len((await db.get_all_puzzle_hashes())) == 2000
assert await db.get_last_derivation_path() == 999
assert await db.get_unused_derivation_path() == 0
assert await db.get_derivation_record(0, 2) == derivation_recs[1]
# Indeces up to 250
await db.set_used_up_to(249)
assert await db.get_unused_derivation_path() == 250
except Exception as e:
print(e, type(e))
await db._clear_database()
await db.close()
db_filename.unlink()
raise e
await db._clear_database()
await db.close()
db_filename.unlink()