diff --git a/chia/_tests/farmer_harvester/test_third_party_harvesters.py b/chia/_tests/farmer_harvester/test_third_party_harvesters.py index df39739aea76..08d8ffdccc76 100644 --- a/chia/_tests/farmer_harvester/test_third_party_harvesters.py +++ b/chia/_tests/farmer_harvester/test_third_party_harvesters.py @@ -7,7 +7,7 @@ import json import logging from os.path import dirname -from typing import Any, Optional, Union, cast +from typing import Optional, Union, cast import pytest from chia_rs import G1Element @@ -25,7 +25,6 @@ from chia.harvester.harvester import Harvester from chia.harvester.harvester_api import HarvesterAPI from chia.protocols import farmer_protocol, full_node_protocol, harvester_protocol, timelord_protocol -from chia.protocols.farmer_protocol import RequestSignedValues from chia.protocols.harvester_protocol import ProofOfSpaceFeeInfo, RespondSignatures, SigningDataKind from chia.protocols.protocol_message_types import ProtocolMessageTypes from chia.server.outbound_message import Message, NodeType, make_msg @@ -115,8 +114,9 @@ async def test_harvester_receive_source_signing_data( finished_validating_data = False farmer_reward_address = decode_puzzle_hash("txch1psqeaw0h244v5sy2r4se8pheyl62n8778zl6t5e7dep0xch9xfkqhx2mej") - async def intercept_harvester_request_signatures(*args: Any) -> Message: - request: harvester_protocol.RequestSignatures = harvester_protocol.RequestSignatures.from_bytes(args[0]) + async def intercept_harvester_request_signatures( + self: HarvesterAPI, request: harvester_protocol.RequestSignatures + ) -> Optional[Message]: nonlocal harvester nonlocal farmer_reward_address @@ -210,24 +210,23 @@ def validate_harvester_request_signatures(request: harvester_protocol.RequestSig data_hash = data.get_hash() assert data_hash == hash - async def intercept_farmer_new_proof_of_space(*args: Any) -> None: + async def intercept_farmer_new_proof_of_space( + self: HarvesterAPI, request: harvester_protocol.NewProofOfSpace, peer: WSChiaConnection + ) -> None: nonlocal farmer nonlocal farmer_reward_address - request: harvester_protocol.NewProofOfSpace = dataclasses.replace( - harvester_protocol.NewProofOfSpace.from_bytes(args[0]), farmer_reward_address_override=farmer_reward_address - ) - peer: WSChiaConnection = args[1] + request = dataclasses.replace(request, farmer_reward_address_override=farmer_reward_address) await FarmerAPI.new_proof_of_space(farmer.server.api, request, peer) - async def intercept_farmer_request_signed_values(*args: Any) -> Optional[Message]: + async def intercept_farmer_request_signed_values( + self: FarmerAPI, request: farmer_protocol.RequestSignedValues + ) -> Optional[Message]: nonlocal farmer nonlocal farmer_reward_address nonlocal full_node_2 - request: RequestSignedValues = RequestSignedValues.from_bytes(args[0]) - # Ensure the FullNode included the source data for the signatures assert request.foliage_block_data assert request.foliage_block_data.get_hash() == request.foliage_block_data_hash @@ -255,6 +254,7 @@ async def intercept_farmer_request_signed_values(*args: Any) -> Optional[Message api=farmer.server.api, handler=intercept_farmer_new_proof_of_space, request_type=ProtocolMessageTypes.new_proof_of_space, + peer_required=True, ) ) exit_stack.enter_context(