| 31 | # message is received from the node-under-test. Subclass P2PInterface and |
| 32 | # override the on_*() methods if you need custom behaviour. |
| 33 | class BaseNode(P2PInterface): |
| 34 | def __init__(self): |
| 35 | """Initialize the P2PInterface |
| 36 | |
| 37 | Used to initialize custom properties for the Node that aren't |
| 38 | included by default in the base class. Be aware that the P2PInterface |
| 39 | base class already stores a counter for each P2P message type and the |
| 40 | last received message of each type, which should be sufficient for the |
| 41 | needs of most tests. |
| 42 | |
| 43 | Call super().__init__() first for standard initialization and then |
| 44 | initialize custom properties.""" |
| 45 | super().__init__() |
| 46 | # Stores a dictionary of all blocks received |
| 47 | self.block_receive_map = defaultdict(int) |
| 48 | |
| 49 | def on_block(self, message): |
| 50 | """Override the standard on_block callback |
| 51 | |
| 52 | Store the hash of a received block in the dictionary.""" |
| 53 | message.block.calc_sha256() |
| 54 | self.block_receive_map[message.block.sha256] += 1 |
| 55 | |
| 56 | def on_inv(self, message): |
| 57 | """Override the standard on_inv callback""" |
| 58 | pass |
| 59 | |
| 60 | def custom_function(): |
| 61 | """Do some custom behaviour |