mypy annotation for test_dos.py (#19036)

* remove unused function, get_block_path()

* add type annotation to test_dos.py

chia/_tests/core/full_node/test_performance.py

@@ -14,7 +14,6 @@
 from chia._tests.util.time_out_assert import time_out_assert
 from chia.consensus.block_record import BlockRecord
 from chia.consensus.pot_iterations import is_overflow_block
-from chia.full_node.full_node_api import FullNodeAPI
 from chia.protocols import full_node_protocol as fnp
 from chia.types.condition_opcodes import ConditionOpcode
 from chia.types.condition_with_args import ConditionWithArgs
@@ -23,16 +22,6 @@
 log = logging.getLogger(__name__)
 
 
-async def get_block_path(full_node: FullNodeAPI):
-    blocks_list = [await full_node.full_node.blockchain.get_full_peak()]
-    assert blocks_list[0] is not None
-    while blocks_list[0].height != 0:
-        b = await full_node.full_node.block_store.get_full_block(blocks_list[0].prev_header_hash)
-        assert b is not None
-        blocks_list.insert(0, b)
-    return blocks_list
-
-
 class TestPerformance:
     @pytest.mark.anyio
     async def test_full_block_performance(

chia/_tests/core/server/test_dos.py

@@ -7,14 +7,14 @@
 
 import pytest
 from aiohttp import ClientSession, ClientTimeout, WSCloseCode, WSMessage, WSMsgType, WSServerHandshakeError
-from chia_rs.sized_ints import uint64
+from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint8, uint16, uint64
 
 import chia.server.server
 from chia._tests.util.time_out_assert import time_out_assert
-from chia.full_node.full_node_api import FullNodeAPI
 from chia.protocols import full_node_protocol
 from chia.protocols.protocol_message_types import ProtocolMessageTypes
-from chia.protocols.shared_protocol import Handshake
+from chia.protocols.shared_protocol import Capability, Handshake
 from chia.server.outbound_message import Message, make_msg
 from chia.server.rate_limits import RateLimiter
 from chia.server.server import ChiaServer
@@ -33,18 +33,10 @@ def not_localhost(host: str) -> bool:
     return False
 
 
-async def get_block_path(full_node: FullNodeAPI):
-    blocks_list = [await full_node.full_node.blockchain.get_full_peak()]
-    assert blocks_list[0] is not None
-    while blocks_list[0].height != 0:
-        b = await full_node.full_node.block_store.get_full_block(blocks_list[0].prev_header_hash)
-        assert b is not None
-        blocks_list.insert(0, b)
-    return blocks_list
-
-
 class FakeRateLimiter:
-    def process_msg_and_check(self, msg, capa, capb) -> Optional[str]:
+    def process_msg_and_check(
+        self, message: Message, our_capabilities: list[Capability], peer_capabilities: list[Capability]
+    ) -> Optional[str]:
         return None
 
 
@@ -142,7 +134,11 @@ async def test_bad_handshake_and_ban(
         await ws.close()
 
     @pytest.mark.anyio
-    async def test_invalid_protocol_handshake(self, setup_two_nodes_fixture, self_hostname):
+    async def test_invalid_protocol_handshake(
+        self,
+        setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
+        self_hostname: str,
+    ) -> None:
         nodes, _, _ = setup_two_nodes_fixture
         server_1 = nodes[0].full_node.server
         server_2 = nodes[1].full_node.server
@@ -159,8 +155,8 @@ async def test_invalid_protocol_handshake(self, setup_two_nodes_fixture, self_ho
         )
 
         # Construct an otherwise valid handshake message
-        handshake: Handshake = Handshake("test", "0.0.32", "1.0.0.0", 3456, 1, [(1, "1")])
-        outbound_handshake: Message = Message(2, None, bytes(handshake))  # 2 is an invalid ProtocolType
+        handshake: Handshake = Handshake("test", "0.0.32", "1.0.0.0", uint16(3456), uint8(1), [(uint16(1), "1")])
+        outbound_handshake: Message = Message(uint8(2), None, bytes(handshake))  # 2 is an invalid ProtocolType
         await ws.send_bytes(bytes(outbound_handshake))
 
         response: WSMessage = await ws.receive()
@@ -173,7 +169,11 @@ async def test_invalid_protocol_handshake(self, setup_two_nodes_fixture, self_ho
         await asyncio.sleep(1)  # give some time for cleanup to work
 
     @pytest.mark.anyio
-    async def test_spam_tx(self, setup_two_nodes_fixture, self_hostname):
+    async def test_spam_tx(
+        self,
+        setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
+        self_hostname: str,
+    ) -> None:
         nodes, _, _ = setup_two_nodes_fixture
         _full_node_1, full_node_2 = nodes
         server_1 = nodes[0].full_node.server
@@ -191,7 +191,7 @@ async def test_spam_tx(self, setup_two_nodes_fixture, self_hostname):
 
         new_tx_message = make_msg(
             ProtocolMessageTypes.new_transaction,
-            full_node_protocol.NewTransaction(bytes([9] * 32), uint64(0), uint64(0)),
+            full_node_protocol.NewTransaction(bytes32([9] * 32), uint64(0), uint64(0)),
         )
         for i in range(4000):
             await ws_con._send_message(new_tx_message)
@@ -215,20 +215,24 @@ async def test_spam_tx(self, setup_two_nodes_fixture, self_hostname):
                 await asyncio.sleep(0)
         await asyncio.sleep(1)
 
-        def is_closed():
+        def is_closed() -> bool:
             return ws_con.closed
 
         await time_out_assert(15, is_closed)
 
         assert ws_con.closed
 
-        def is_banned():
+        def is_banned() -> bool:
             return "1.2.3.4" in server_2.banned_peers
 
         await time_out_assert(15, is_banned)
 
     @pytest.mark.anyio
-    async def test_spam_message_non_tx(self, setup_two_nodes_fixture, self_hostname):
+    async def test_spam_message_non_tx(
+        self,
+        setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
+        self_hostname: str,
+    ) -> None:
         nodes, _, _ = setup_two_nodes_fixture
         _full_node_1, full_node_2 = nodes
         server_1 = nodes[0].full_node.server
@@ -244,7 +248,7 @@ async def test_spam_message_non_tx(self, setup_two_nodes_fixture, self_hostname)
         ws_con.peer_info = PeerInfo("1.2.3.4", ws_con.peer_info.port)
         ws_con_2.peer_info = PeerInfo("1.2.3.4", ws_con_2.peer_info.port)
 
-        def is_closed():
+        def is_closed() -> bool:
             return ws_con.closed
 
         new_message = make_msg(
@@ -271,13 +275,17 @@ def is_closed():
         await time_out_assert(15, is_closed)
 
         # Banned
-        def is_banned():
+        def is_banned() -> bool:
             return "1.2.3.4" in server_2.banned_peers
 
         await time_out_assert(15, is_banned)
 
     @pytest.mark.anyio
-    async def test_spam_message_too_large(self, setup_two_nodes_fixture, self_hostname):
+    async def test_spam_message_too_large(
+        self,
+        setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
+        self_hostname: str,
+    ) -> None:
         nodes, _, _ = setup_two_nodes_fixture
         _full_node_1, full_node_2 = nodes
         server_1 = nodes[0].full_node.server
@@ -293,7 +301,7 @@ async def test_spam_message_too_large(self, setup_two_nodes_fixture, self_hostna
         ws_con.peer_info = PeerInfo("1.2.3.4", ws_con.peer_info.port)
         ws_con_2.peer_info = PeerInfo("1.2.3.4", ws_con_2.peer_info.port)
 
-        def is_closed():
+        def is_closed() -> bool:
             return ws_con.closed
 
         new_message = make_msg(
@@ -307,13 +315,13 @@ def is_closed():
         assert not ws_con.closed
 
         # Remove outbound rate limiter to test inbound limits
-        ws_con.outbound_rate_limiter = FakeRateLimiter()
+        ws_con.outbound_rate_limiter = FakeRateLimiter()  # type: ignore[assignment]
 
         await ws_con._send_message(new_message)
         await time_out_assert(15, is_closed)
 
         # Banned
-        def is_banned():
+        def is_banned() -> bool:
             return "1.2.3.4" in server_2.banned_peers
 
         await time_out_assert(15, is_banned)

mypy-exclusions.txt

@@ -59,7 +59,6 @@ chia._tests.core.full_node.test_address_manager
 chia._tests.core.full_node.test_node_load
 chia._tests.core.full_node.test_performance
 chia._tests.core.full_node.test_transactions
-chia._tests.core.server.test_dos
 chia._tests.core.server.test_rate_limits
 chia._tests.core.ssl.test_ssl
 chia._tests.core.test_crawler_rpc