chia-blockchain/tests/setup_nodes.py

67 lines
2.1 KiB
Python

import os
from typing import Any, Dict
from src.blockchain import Blockchain
from src.store import FullNodeStore
from src.full_node import FullNode
from src.server.connection import NodeType
from src.server.server import ChiaServer
from src.types.full_block import FullBlock
from tests.block_tools import BlockTools
bt = BlockTools()
test_constants: Dict[str, Any] = {
"DIFFICULTY_STARTING": 1,
"DISCRIMINANT_SIZE_BITS": 16,
"BLOCK_TIME_TARGET": 10,
"MIN_BLOCK_TIME": 2,
"DIFFICULTY_FACTOR": 3,
"DIFFICULTY_EPOCH": 12, # The number of blocks per epoch
"DIFFICULTY_WARP_FACTOR": 4, # DELAY divides EPOCH in order to warp efficiently.
"DIFFICULTY_DELAY": 3, # EPOCH / WARP_FACTOR
}
test_constants["GENESIS_BLOCK"] = bytes(
bt.create_genesis_block(test_constants, bytes([0] * 32), b"0")
)
async def setup_two_nodes():
"""
Setup and teardown of two full nodes, with blockchains and separate DBs.
"""
# SETUP
store_1 = await FullNodeStore.create("blockchain_test")
store_2 = await FullNodeStore.create("blockchain_test_2")
await store_1._clear_database()
await store_2._clear_database()
b_1: Blockchain = await Blockchain.create({}, test_constants)
b_2: Blockchain = await Blockchain.create({}, test_constants)
await store_1.add_block(FullBlock.from_bytes(test_constants["GENESIS_BLOCK"]))
await store_2.add_block(FullBlock.from_bytes(test_constants["GENESIS_BLOCK"]))
full_node_1 = FullNode(store_1, b_1)
server_1 = ChiaServer(21234, full_node_1, NodeType.FULL_NODE)
_ = await server_1.start_server("127.0.0.1", full_node_1._on_connect)
full_node_1._set_server(server_1)
full_node_2 = FullNode(store_2, b_2)
server_2 = ChiaServer(21235, full_node_2, NodeType.FULL_NODE)
full_node_2._set_server(server_2)
yield (full_node_1, full_node_2, server_1, server_2)
# TEARDOWN
full_node_1._shutdown()
full_node_2._shutdown()
server_1.close_all()
server_2.close_all()
await server_1.await_closed()
await server_2.await_closed()
await store_1.close()
await store_2.close()
os.remove("blockchain_test")
os.remove("blockchain_test_2")