Repository: Chia-Network/chia-blockchain
File: chia/_tests/core/full_node/test_full_node.py (Python)

Code before:

@pytest.mark.anyio
@pytest.mark.parametrize("mismatch_cost", [True, False])
@pytest.mark.parametrize("mismatch_fee", [True, False])
async def test_ban_for_mismatched_tx_cost_fee(
    setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
    self_hostname: str,
    mismatch_cost: bool,
    mismatch_fee: bool,
) -> None:
    """
    Tests that a peer gets banned if it sends a `NewTransaction` message with a
    cost and/or fee that doesn't match the transaction's validation cost/fee.
    We setup two full nodes with the test transaction as already seen, and we
    check its validation cost and fee against the ones specified in the
    `NewTransaction` message.
    """

    async def generate_and_send_tx(node_to_add_tx_to: FullNode, ws_con: WSChiaConnection) -> MempoolItem:
        wallet = WalletTool(test_constants)
        wallet_ph = wallet.get_new_puzzlehash()
        blocks = bt.get_consecutive_blocks(
            3, guarantee_transaction_block=True, farmer_reward_puzzle_hash=wallet_ph, pool_reward_puzzle_hash=wallet_ph
        )
        for block in blocks:
            await node_to_add_tx_to.add_block(block)
        coin = blocks[-1].get_included_reward_coins()[0]
        spend_bundle = wallet.generate_signed_transaction(uint64(42), wallet_ph, coin)
        sb_name = spend_bundle.name()
        await node_to_add_tx_to.add_transaction(spend_bundle, sb_name, ws_con)
        mempool_item = node_to_add_tx_to.mempool_manager.get_mempool_item(sb_name)
        assert mempool_item is not None
        return mempool_item

    nodes, _, bt = setup_two_nodes_fixture
    full_node_1, full_node_2 = nodes
    server_1 = full_node_1.full_node.server
    server_2 = full_node_2.full_node.server
    await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), full_node_2.full_node.on_connect)
    ws_con_1 = next(iter(server_1.all_connections.values()))
    ws_con_2 = next(iter(server_2.all_connections.values()))
    # Create a transaction and add it to the relevant full node's mempool
    mempool_item = await generate_and_send_tx(full_node_1.full_node, ws_con_1)
    # Now send a NewTransaction with a cost and/or fee mismatch from the second
    # full node.
    cost = uint64(mempool_item.cost + 1) if mismatch_cost else mempool_item.cost
    fee = uint64(mempool_item.fee + 1) if mismatch_fee else mempool_item.fee
    msg = make_msg(ProtocolMessageTypes.new_transaction, NewTransaction(mempool_item.name, cost, fee))
    # We won't ban localhost, so let's set a different ip address
    full_node_2_ip = "1.3.3.7"
    ws_con_1.peer_info = PeerInfo(full_node_2_ip, ws_con_1.peer_info.port)

Review comment (suggestion):
the name of this function, `generate_and_send_tx()` seems misleading. You also generate blocks and add them to the chain.
It's also only called from one place, so there might be a better line to cut along to pull out a function

Code after:
@pytest.mark.anyio
@pytest.mark.parametrize("mismatch_cost", [True, False])
@pytest.mark.parametrize("mismatch_fee", [True, False])
async def test_ban_for_mismatched_tx_cost_fee(
    setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
    self_hostname: str,
    mismatch_cost: bool,
    mismatch_fee: bool,
) -> None:
    """
    Tests that a peer gets banned if it sends a `NewTransaction` message with a
    cost and/or fee that doesn't match the transaction's validation cost/fee.
    We setup two full nodes with the test transaction as already seen, and we
    check its validation cost and fee against the ones specified in the
    `NewTransaction` message.
    """
    nodes, _, bt = setup_two_nodes_fixture
    full_node_1, full_node_2 = nodes
    server_1 = full_node_1.full_node.server
    server_2 = full_node_2.full_node.server
    await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), full_node_2.full_node.on_connect)
    ws_con_1 = next(iter(server_1.all_connections.values()))
    ws_con_2 = next(iter(server_2.all_connections.values()))
    wallet = WalletTool(test_constants)
    wallet_ph = wallet.get_new_puzzlehash()
    blocks = bt.get_consecutive_blocks(
        3, guarantee_transaction_block=True, farmer_reward_puzzle_hash=wallet_ph, pool_reward_puzzle_hash=wallet_ph
    )
    for block in blocks:
        await full_node_1.full_node.add_block(block)
    # Create a transaction and add it to the relevant full node's mempool
    coin = blocks[-1].get_included_reward_coins()[0]
    sb = wallet.generate_signed_transaction(uint64(42), wallet_ph, coin)
    sb_name = sb.name()
    await full_node_1.full_node.add_transaction(sb, sb_name, ws_con_1)
    mempool_item = full_node_1.full_node.mempool_manager.get_mempool_item(sb_name)
    assert mempool_item is not None
    # Now send a NewTransaction with a cost and/or fee mismatch from the second
    # full node.
    cost = uint64(mempool_item.cost + 1) if mismatch_cost else mempool_item.cost
    fee = uint64(mempool_item.fee + 1) if mismatch_fee else mempool_item.fee
    msg = make_msg(ProtocolMessageTypes.new_transaction, NewTransaction(mempool_item.name, cost, fee))
    # We won't ban localhost, so let's set a different ip address for the
    # second node.
    full_node_2_ip = "1.3.3.7"
    ws_con_1.peer_info = PeerInfo(full_node_2_ip, ws_con_1.peer_info.port)
    # Send the NewTransaction message from the second node to the first
    await ws_con_2.send_message(msg)
    # Make sure the first full node has banned the second as the item it has
    # already seen has a different validation cost and/or fee than the one from
    # the NewTransaction message.