From 60713c9f91c888ebaa61f7f20a78ea474a1d07a5 Mon Sep 17 00:00:00 2001
From: Michal Nazarewicz <mina86@mina86.com>
Date: Thu, 1 Jul 2021 22:04:05 +0200
Subject: [PATCH] pytest: make run_handshake handle unsolicited Block messages
 (#4437)

This addresses the following failure:

    Traceback (most recent call last):
      File "tests/spec/network/peers_request.py", line 44, in <module>
        asyncio.run(main())
      File "/usr/lib/python3.7/asyncio/runners.py", line 43, in run
        return loop.run_until_complete(main)
      File "/usr/lib/python3.7/asyncio/base_events.py", line 579, in run_until_complete
        return future.result()
      File "tests/spec/network/peers_request.py", line 35, in main
        await run_handshake(conn1, nodes[0].node_key.pk, key_pair_1, listen_port=12346)
      File "lib/peer.py", line 182, in run_handshake
        assert response.enum == 'Handshake', response.enum if response.enum != 'HandshakeFailure' else response.HandshakeFailure[1].enum
    AssertionError: Block

Fixes: https://github.com/near/nearcore/issues/4432
---
 pytest/lib/peer.py | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/pytest/lib/peer.py b/pytest/lib/peer.py
index 20b4e6bb6ba..465baefd05f 100644
--- a/pytest/lib/peer.py
+++ b/pytest/lib/peer.py
@@ -161,25 +161,29 @@ async def run_handshake(conn: Connection,
                         key_pair: SigningKey,
                         listen_port=12345):
     handshake = create_handshake(key_pair, target_public_key, listen_port)
-    sign_handshake(key_pair, handshake.Handshake)
 
-    await conn.send(handshake)
-    response = await conn.recv()
+    async def send_handshake():
+        sign_handshake(key_pair, handshake.Handshake)
+        await conn.send(handshake)
+        # The peer might sent us an unsolicited message before replying to
+        # a successful handshake.  This is because node is multi-threaded and
+        # peers are added to PeerManager before the reply is sent.  Since we
+        # don’t care about those messages, ignore them and wait for some kind of
+        # Handshake reply.
+        return await conn.recv(lambda msg: msg.enum.startswith('Handshake'))
+
+    response = await send_handshake()
 
     if response.enum == 'HandshakeFailure' and response.HandshakeFailure[1].enum == 'ProtocolVersionMismatch':
         pvm = response.HandshakeFailure[1].ProtocolVersionMismatch.version
         handshake.Handshake.version = pvm
-        sign_handshake(key_pair, handshake.Handshake)
-        await conn.send(handshake)
-        response = await conn.recv()
+        response = await send_handshake()
 
     if response.enum == 'HandshakeFailure' and response.HandshakeFailure[1].enum == 'GenesisMismatch':
         gm = response.HandshakeFailure[1].GenesisMismatch
         handshake.Handshake.chain_info.genesis_id.chain_id = gm.chain_id
         handshake.Handshake.chain_info.genesis_id.hash = gm.hash
-        sign_handshake(key_pair, handshake.Handshake)
-        await conn.send(handshake)
-        response = await conn.recv()
+        response = await send_handshake()
 
     assert response.enum == 'Handshake', response.enum if response.enum != 'HandshakeFailure' else response.HandshakeFailure[1].enum