8080 Expect: disconnect.
8181'''
8282
83- class BaseNode (NodeConnCB ):
83+ direct_fetch_response_time = 0.05
84+
85+ class BaseNode (SingleNodeConnCB ):
8486 def __init__ (self ):
85- NodeConnCB .__init__ (self )
86- self .connection = None
87+ SingleNodeConnCB .__init__ (self )
8788 self .last_inv = None
8889 self .last_headers = None
8990 self .last_block = None
90- self .ping_counter = 1
91- self .last_pong = msg_pong (0 )
9291 self .last_getdata = None
93- self .sleep_time = 0.05
9492 self .block_announced = False
9593 self .last_getheaders = None
9694 self .disconnected = False
95+ self .last_blockhash_announced = None
9796
9897 def clear_last_announcement (self ):
9998 with mininode_lock :
10099 self .block_announced = False
101100 self .last_inv = None
102101 self .last_headers = None
103102
104- def add_connection (self , conn ):
105- self .connection = conn
106-
107103 # Request data for a list of block hashes
108104 def get_data (self , block_hashes ):
109105 msg = msg_getdata ()
@@ -122,17 +118,17 @@ def send_block_inv(self, blockhash):
122118 msg .inv = [CInv (2 , blockhash )]
123119 self .connection .send_message (msg )
124120
125- # Wrapper for the NodeConn's send_message function
126- def send_message (self , message ):
127- self .connection .send_message (message )
128-
129121 def on_inv (self , conn , message ):
130122 self .last_inv = message
131123 self .block_announced = True
124+ self .last_blockhash_announced = message .inv [- 1 ].hash
132125
133126 def on_headers (self , conn , message ):
134127 self .last_headers = message
135- self .block_announced = True
128+ if len (message .headers ):
129+ self .block_announced = True
130+ message .headers [- 1 ].calc_sha256 ()
131+ self .last_blockhash_announced = message .headers [- 1 ].sha256
136132
137133 def on_block (self , conn , message ):
138134 self .last_block = message .block
@@ -141,9 +137,6 @@ def on_block(self, conn, message):
141137 def on_getdata (self , conn , message ):
142138 self .last_getdata = message
143139
144- def on_pong (self , conn , message ):
145- self .last_pong = message
146-
147140 def on_getheaders (self , conn , message ):
148141 self .last_getheaders = message
149142
@@ -157,7 +150,7 @@ def check_last_announcement(self, headers=None, inv=None):
157150 expect_headers = headers if headers != None else []
158151 expect_inv = inv if inv != None else []
159152 test_function = lambda : self .block_announced
160- self . sync ( test_function )
153+ assert ( wait_until ( test_function , timeout = 60 ) )
161154 with mininode_lock :
162155 self .block_announced = False
163156
@@ -180,43 +173,32 @@ def check_last_announcement(self, headers=None, inv=None):
180173 return success
181174
182175 # Syncing helpers
183- def sync (self , test_function , timeout = 60 ):
184- while timeout > 0 :
185- with mininode_lock :
186- if test_function ():
187- return
188- time .sleep (self .sleep_time )
189- timeout -= self .sleep_time
190- raise AssertionError ("Sync failed to complete" )
191-
192- def sync_with_ping (self , timeout = 60 ):
193- self .send_message (msg_ping (nonce = self .ping_counter ))
194- test_function = lambda : self .last_pong .nonce == self .ping_counter
195- self .sync (test_function , timeout )
196- self .ping_counter += 1
197- return
198-
199176 def wait_for_block (self , blockhash , timeout = 60 ):
200177 test_function = lambda : self .last_block != None and self .last_block .sha256 == blockhash
201- self . sync ( test_function , timeout )
178+ assert ( wait_until ( test_function , timeout = timeout ) )
202179 return
203180
204181 def wait_for_getheaders (self , timeout = 60 ):
205182 test_function = lambda : self .last_getheaders != None
206- self . sync ( test_function , timeout )
183+ assert ( wait_until ( test_function , timeout = timeout ) )
207184 return
208185
209186 def wait_for_getdata (self , hash_list , timeout = 60 ):
210187 if hash_list == []:
211188 return
212189
213190 test_function = lambda : self .last_getdata != None and [x .hash for x in self .last_getdata .inv ] == hash_list
214- self . sync ( test_function , timeout )
191+ assert ( wait_until ( test_function , timeout = timeout ) )
215192 return
216193
217194 def wait_for_disconnect (self , timeout = 60 ):
218195 test_function = lambda : self .disconnected
219- self .sync (test_function , timeout )
196+ assert (wait_until (test_function , timeout = timeout ))
197+ return
198+
199+ def wait_for_block_announcement (self , block_hash , timeout = 60 ):
200+ test_function = lambda : self .last_blockhash_announced == block_hash
201+ assert (wait_until (test_function , timeout = timeout ))
220202 return
221203
222204 def send_header_for_blocks (self , new_blocks ):
@@ -266,7 +248,9 @@ def mine_blocks(self, count):
266248 def mine_reorg (self , length ):
267249 self .nodes [0 ].generate (length ) # make sure all invalidated blocks are node0's
268250 sync_blocks (self .nodes , wait = 0.1 )
269- [x .clear_last_announcement () for x in self .p2p_connections ]
251+ for x in self .p2p_connections :
252+ x .wait_for_block_announcement (int (self .nodes [0 ].getbestblockhash (), 16 ))
253+ x .clear_last_announcement ()
270254
271255 tip_height = self .nodes [1 ].getblockcount ()
272256 hash_to_invalidate = self .nodes [1 ].getblockhash (tip_height - (length - 1 ))
@@ -495,7 +479,7 @@ def run_test(self):
495479
496480 test_node .send_header_for_blocks (blocks )
497481 test_node .sync_with_ping ()
498- test_node .wait_for_getdata ([x .sha256 for x in blocks ], timeout = test_node . sleep_time )
482+ test_node .wait_for_getdata ([x .sha256 for x in blocks ], timeout = direct_fetch_response_time )
499483
500484 [ test_node .send_message (msg_block (x )) for x in blocks ]
501485
@@ -526,13 +510,13 @@ def run_test(self):
526510 # both blocks (same work as tip)
527511 test_node .send_header_for_blocks (blocks [1 :2 ])
528512 test_node .sync_with_ping ()
529- test_node .wait_for_getdata ([x .sha256 for x in blocks [0 :2 ]], timeout = test_node . sleep_time )
513+ test_node .wait_for_getdata ([x .sha256 for x in blocks [0 :2 ]], timeout = direct_fetch_response_time )
530514
531515 # Announcing 16 more headers should trigger direct fetch for 14 more
532516 # blocks
533517 test_node .send_header_for_blocks (blocks [2 :18 ])
534518 test_node .sync_with_ping ()
535- test_node .wait_for_getdata ([x .sha256 for x in blocks [2 :16 ]], timeout = test_node . sleep_time )
519+ test_node .wait_for_getdata ([x .sha256 for x in blocks [2 :16 ]], timeout = direct_fetch_response_time )
536520
537521 # Announcing 1 more header should not trigger any response
538522 test_node .last_getdata = None
0 commit comments