Differential D9341 Diff 28019 test/functional/test_framework/mininode.py

Changeset View

Standalone View

test/functional/test_framework/mininode.py

Show All 17 Lines

import asyncio		import asyncio
from collections import defaultdict		from collections import defaultdict
from io import BytesIO		from io import BytesIO
import logging		import logging
import struct		import struct
import sys		import sys
import threading		import threading
		from typing import List

from test_framework.messages import (		from test_framework.messages import (
		AvalancheResponse,
CBlockHeader,		CBlockHeader,
		CInv,
MIN_VERSION_SUPPORTED,		MIN_VERSION_SUPPORTED,
msg_addr,		msg_addr,
msg_addrv2,		msg_addrv2,
		msg_avahello,
msg_avapoll,		msg_avapoll,
msg_avaproof,		msg_avaproof,
msg_tcpavaresponse,		msg_tcpavaresponse,
msg_avahello,
msg_block,		msg_block,
MSG_BLOCK,		MSG_BLOCK,
msg_blocktxn,		msg_blocktxn,
msg_cfcheckpt,		msg_cfcheckpt,
msg_cfheaders,		msg_cfheaders,
msg_cfilter,		msg_cfilter,
msg_cmpctblock,		msg_cmpctblock,
msg_feefilter,		msg_feefilter,
Show All 17 Lines	from test_framework.messages import (
msg_sendheaders,		msg_sendheaders,
msg_tx,		msg_tx,
MSG_TX,		MSG_TX,
MSG_TYPE_MASK,		MSG_TYPE_MASK,
msg_verack,		msg_verack,
msg_version,		msg_version,
NODE_NETWORK,		NODE_NETWORK,
sha256,		sha256,
		TCPAvalancheResponse,
)		)
from test_framework.util import wait_until		from test_framework.util import wait_until

logger = logging.getLogger("TestFramework.mininode")		logger = logging.getLogger("TestFramework.mininode")

MESSAGEMAP = {		MESSAGEMAP = {
b"addr": msg_addr,		b"addr": msg_addr,
b"addrv2": msg_addrv2,		b"addrv2": msg_addrv2,
▲ Show 20 Lines • Show All 709 Lines • ▼ Show 20 Lines	def wait_for_broadcast(self, txns, timeout=60):
"""Waits for the txns (list of txids) to complete initial broadcast.		"""Waits for the txns (list of txids) to complete initial broadcast.
The mempool should mark unbroadcast=False for these transactions.		The mempool should mark unbroadcast=False for these transactions.
"""		"""
# Wait until invs have been received (and getdatas sent) for each txid.		# Wait until invs have been received (and getdatas sent) for each txid.
self.wait_until(lambda: set(self.get_invs()) == set(		self.wait_until(lambda: set(self.get_invs()) == set(
[int(tx, 16) for tx in txns]), timeout)		[int(tx, 16) for tx in txns]), timeout)
# Flush messages and wait for the getdatas to be processed		# Flush messages and wait for the getdatas to be processed
self.sync_with_ping()		self.sync_with_ping()


		class P2PAvalancheTestNode(P2PInterface):
		"""A P2PInterface supporting Avalanche messages"""
		DUMMY_PROOFID = 1337

		def __init__(self):
		self.round = 0
		self.avahello = None
		self.avaresponses = []
		self.avapolls = []
		self.avaproof = None
		super().__init__()

		def peer_connect(self, args, *kwargs):
		create_conn = super().peer_connect(args, *kwargs)

		# Save the nonce and extra entropy so they can be reused later.
		self.local_nonce = self.on_connection_send_msg.nNonce
		self.local_extra_entropy = self.on_connection_send_msg.nExtraEntropy

		return create_conn

		def on_version(self, message):
		super().on_version(message)

		# Save the nonce and extra entropy so they can be reused later.
		self.remote_nonce = message.nNonce
		self.remote_extra_entropy = message.nExtraEntropy

		def on_avaresponse(self, message):
		with mininode_lock:
		self.avaresponses.append(message.response)

		def on_avapoll(self, message):
		with mininode_lock:
		self.avapolls.append(message.poll)

		def on_avahello(self, message):
		with mininode_lock:
		assert(self.avahello is None)
		self.avahello = message

		def send_avaresponse(self, round, votes, privkey):
		response = AvalancheResponse(round, 0, votes)
		sig = privkey.sign_schnorr(response.get_hash())
		msg = msg_tcpavaresponse()
		msg.response = TCPAvalancheResponse(response, sig)
		self.send_message(msg)

		def wait_for_avaresponse(self, timeout=5):
		wait_until(
		lambda: len(self.avaresponses) > 0,
		timeout=timeout,
		lock=mininode_lock)

		with mininode_lock:
		return self.avaresponses.pop(0)

		def send_poll(self, hashes):
		msg = msg_avapoll()
		msg.poll.round = self.round
		self.round += 1
		for h in hashes:
		msg.poll.invs.append(CInv(2, h))
		self.send_message(msg)

		def get_avapoll_if_available(self):
		with mininode_lock:
		return self.avapolls.pop(0) if len(self.avapolls) > 0 else None

		def wait_for_avahello(self, timeout=5):
		wait_until(
		lambda: self.avahello is not None,
		timeout=timeout,
		lock=mininode_lock)

		with mininode_lock:
		return self.avahello

		def send_avahello(self):
		msg = msg_avahello()
		msg.hello.delegation.proofid = self.DUMMY_PROOFID
		self.send_message(msg)

		def send_getdata(self, inv: List[CInv]):
		msg = msg_getdata()
		msg.inv = inv
		self.send_message(msg)

		def on_avaproof(self, message):
		with mininode_lock:
		assert(self.avaproof is None)
		self.avaproof = message

		def wait_for_avaproof(self, timeout=10):
		wait_until(
		lambda: self.avaproof is not None,
		timeout=timeout,
		lock=mininode_lock)

		with mininode_lock:
		return self.avaproof