[qa] mininode: Expose connection state through is_connected

MarcoFalke · MarcoFalke · commit fa1eac9cdb1a · 2018-06-22T12:00:00.000-04:00
diff --git a/test/functional/feature_assumevalid.py b/test/functional/feature_assumevalid.py
@@ -68,12 +68,12 @@ def setup_network(self):
     def send_blocks_until_disconnected(self, p2p_conn):
         """Keep sending blocks to the node until we're disconnected."""
         for i in range(len(self.blocks)):
-            if p2p_conn.state != "connected":
+            if not p2p_conn.is_connected:
                 break
             try:
                 p2p_conn.send_message(msg_block(self.blocks[i]))
             except IOError as e:
-                assert str(e) == 'Not connected, no pushbuf'
+                assert not p2p_conn.is_connected
                 break
 
     def assert_blockchain_height(self, node, height):
diff --git a/test/functional/p2p_compactblocks.py b/test/functional/p2p_compactblocks.py
@@ -87,7 +87,7 @@ def send_await_disconnect(self, message, timeout=30):
         This is used when we want to send a message into the node that we expect
         will get us disconnected, eg an invalid block."""
         self.send_message(message)
-        wait_until(lambda: self.state != "connected", timeout=timeout, lock=mininode_lock)
+        wait_until(lambda: not self.is_connected, timeout=timeout, lock=mininode_lock)
 
 class CompactBlocksTest(BitcoinTestFramework):
     def set_test_params(self):
diff --git a/test/functional/p2p_leak.py b/test/functional/p2p_leak.py
@@ -118,11 +118,11 @@ def run_test(self):
         time.sleep(5)
 
         #This node should have been banned
-        assert no_version_bannode.state != "connected"
+        assert not no_version_bannode.is_connected
 
         # These nodes should have been disconnected
-        assert unsupported_service_bit5_node.state != "connected"
-        assert unsupported_service_bit7_node.state != "connected"
+        assert not unsupported_service_bit5_node.is_connected
+        assert not unsupported_service_bit7_node.is_connected
 
         self.nodes[0].disconnect_p2ps()
 
diff --git a/test/functional/p2p_timeouts.py b/test/functional/p2p_timeouts.py
@@ -47,9 +47,9 @@ def run_test(self):
 
         sleep(1)
 
-        assert no_verack_node.connected
-        assert no_version_node.connected
-        assert no_send_node.connected
+        assert no_verack_node.is_connected
+        assert no_version_node.is_connected
+        assert no_send_node.is_connected
 
         no_verack_node.send_message(msg_ping())
         no_version_node.send_message(msg_ping())
@@ -58,18 +58,18 @@ def run_test(self):
 
         assert "version" in no_verack_node.last_message
 
-        assert no_verack_node.connected
-        assert no_version_node.connected
-        assert no_send_node.connected
+        assert no_verack_node.is_connected
+        assert no_version_node.is_connected
+        assert no_send_node.is_connected
 
         no_verack_node.send_message(msg_ping())
         no_version_node.send_message(msg_ping())
 
         sleep(31)
 
-        assert not no_verack_node.connected
-        assert not no_version_node.connected
-        assert not no_send_node.connected
+        assert not no_verack_node.is_connected
+        assert not no_version_node.is_connected
+        assert not no_send_node.is_connected
 
 if __name__ == '__main__':
     TimeoutsTest().main()
diff --git a/test/functional/test_framework/mininode.py b/test/functional/test_framework/mininode.py
@@ -77,14 +77,20 @@ def __init__(self):
 
         super().__init__(map=mininode_socket_map)
 
+        self._conn_open = False
+
+    @property
+    def is_connected(self):
+        return self._conn_open
+
     def peer_connect(self, dstaddr, dstport, net="regtest"):
         self.dstaddr = dstaddr
         self.dstport = dstport
         self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
         self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
         self.sendbuf = b""
         self.recvbuf = b""
-        self.state = "connecting"
+        self._asyncore_pre_connection = True
         self.network = net
         self.disconnect = False
 
@@ -97,22 +103,23 @@ def peer_connect(self, dstaddr, dstport, net="regtest"):
 
     def peer_disconnect(self):
         # Connection could have already been closed by other end.
-        if self.state == "connected":
-            self.disconnect_node()
+        if self.is_connected:
+            self.disconnect = True  # Signal asyncore to disconnect
 
     # Connection and disconnection methods
 
     def handle_connect(self):
         """asyncore callback when a connection is opened."""
-        if self.state != "connected":
+        if not self.is_connected:
             logger.debug("Connected & Listening: %s:%d" % (self.dstaddr, self.dstport))
-            self.state = "connected"
+            self._conn_open = True
+            self._asyncore_pre_connection = False
             self.on_open()
 
     def handle_close(self):
         """asyncore callback when a connection is closed."""
         logger.debug("Closing connection to: %s:%d" % (self.dstaddr, self.dstport))
-        self.state = "closed"
+        self._conn_open = False
         self.recvbuf = b""
         self.sendbuf = b""
         try:
@@ -121,13 +128,6 @@ def handle_close(self):
             pass
         self.on_close()
 
-    def disconnect_node(self):
-        """Disconnect the p2p connection.
-
-        Called by the test logic thread. Causes the p2p connection
-        to be disconnected on the next iteration of the asyncore loop."""
-        self.disconnect = True
-
     # Socket read methods
 
     def handle_read(self):
@@ -182,17 +182,16 @@ def on_message(self, message):
     def writable(self):
         """asyncore method to determine whether the handle_write() callback should be called on the next loop."""
         with mininode_lock:
-            pre_connection = self.state == "connecting"
             length = len(self.sendbuf)
-        return (length > 0 or pre_connection)
+        return length > 0 or self._asyncore_pre_connection
 
     def handle_write(self):
         """asyncore callback when data should be written to the socket."""
         with mininode_lock:
             # asyncore does not expose socket connection, only the first read/write
             # event, thus we must check connection manually here to know when we
             # actually connect
-            if self.state == "connecting":
+            if self._asyncore_pre_connection:
                 self.handle_connect()
             if not self.writable():
                 return
@@ -204,26 +203,17 @@ def handle_write(self):
                 return
             self.sendbuf = self.sendbuf[sent:]
 
-    def send_message(self, message, pushbuf=False):
+    def send_message(self, message):
         """Send a P2P message over the socket.
 
         This method takes a P2P payload, builds the P2P header and adds
         the message to the send buffer to be sent over the socket."""
-        if self.state != "connected" and not pushbuf:
-            raise IOError('Not connected, no pushbuf')
+        if not self.is_connected:
+            raise IOError('Not connected')
         self._log_message("send", message)
-        command = message.command
-        data = message.serialize()
-        tmsg = MAGIC_BYTES[self.network]
-        tmsg += command
-        tmsg += b"\x00" * (12 - len(command))
-        tmsg += struct.pack("<I", len(data))
-        th = sha256(data)
-        h = sha256(th)
-        tmsg += h[:4]
-        tmsg += data
+        tmsg = self._build_message(message)
         with mininode_lock:
-            if (len(self.sendbuf) == 0 and not pushbuf):
+            if len(self.sendbuf) == 0:
                 try:
                     sent = self.send(tmsg)
                     self.sendbuf = tmsg[sent:]
@@ -234,6 +224,20 @@ def send_message(self, message, pushbuf=False):
 
     # Class utility methods
 
+    def _build_message(self, message):
+        """Build a serialized P2P message"""
+        command = message.command
+        data = message.serialize()
+        tmsg = MAGIC_BYTES[self.network]
+        tmsg += command
+        tmsg += b"\x00" * (12 - len(command))
+        tmsg += struct.pack("<I", len(data))
+        th = sha256(data)
+        h = sha256(th)
+        tmsg += h[:4]
+        tmsg += data
+        return tmsg
+
     def _log_message(self, direction, msg):
         """Logs a message being sent or received over the connection."""
         if direction == "send":
@@ -280,7 +284,7 @@ def peer_connect(self, *args, services=NODE_NETWORK|NODE_WITNESS, send_version=T
             vt.addrTo.port = self.dstport
             vt.addrFrom.ip = "0.0.0.0"
             vt.addrFrom.port = 0
-            self.send_message(vt, True)
+            self.sendbuf = self._build_message(vt)  # Will be sent right after handle_connect
 
     # Message receiving methods
 
@@ -348,7 +352,7 @@ def on_version(self, message):
     # Connection helper methods
 
     def wait_for_disconnect(self, timeout=60):
-        test_function = lambda: self.state != "connected"
+        test_function = lambda: not self.is_connected
         wait_until(test_function, timeout=timeout, lock=mininode_lock)
 
     # Message receiving helper methods
