import dataclasses import logging import random import time from secrets import token_bytes from typing import Any, Callable, List, Tuple, Type, Union import pytest from blspy import G1Element from chia.plot_sync.delta import Delta from chia.plot_sync.receiver import Receiver, Sync from chia.plot_sync.util import ErrorCodes, State from chia.protocols.harvester_protocol import ( Plot, PlotSyncDone, PlotSyncIdentifier, PlotSyncPathList, PlotSyncPlotList, PlotSyncResponse, PlotSyncStart, ) from chia.server.ws_connection import NodeType from chia.types.blockchain_format.sized_bytes import bytes32 from chia.util.ints import uint8, uint32, uint64 from chia.util.misc import get_list_or_len from chia.util.streamable import _T_Streamable from tests.plot_sync.util import get_dummy_connection log = logging.getLogger(__name__) next_message_id = uint64(0) def assert_default_values(receiver: Receiver) -> None: assert receiver.current_sync() == Sync() assert receiver.last_sync() == Sync() assert receiver.plots() == {} assert receiver.invalid() == [] assert receiver.keys_missing() == [] assert receiver.duplicates() == [] async def dummy_callback(_: bytes32, __: Delta) -> None: pass class SyncStepData: state: State function: Any payload_type: Any args: Any def __init__( self, state: State, function: Callable[[_T_Streamable], Any], payload_type: Type[_T_Streamable], *args: Any ) -> None: self.state = state self.function = function self.payload_type = payload_type self.args = args def plot_sync_identifier(current_sync_id: uint64, message_id: uint64) -> PlotSyncIdentifier: return PlotSyncIdentifier(uint64(0), current_sync_id, message_id) def create_payload(payload_type: Any, start: bool, *args: Any) -> Any: global next_message_id if start: next_message_id = uint64(0) next_identifier = plot_sync_identifier(uint64(1), next_message_id) next_message_id = uint64(next_message_id + 1) return payload_type(next_identifier, *args) def assert_error_response(plot_sync: Receiver, error_code: ErrorCodes) -> None: connection = plot_sync.connection() assert connection is not None message = connection.last_sent_message assert message is not None response: PlotSyncResponse = PlotSyncResponse.from_bytes(message.data) assert response.error is not None assert response.error.code == error_code.value def pre_function_validate(receiver: Receiver, data: Union[List[Plot], List[str]], expected_state: State) -> None: if expected_state == State.loaded: for plot_info in data: assert type(plot_info) == Plot assert plot_info.filename not in receiver.plots() elif expected_state == State.removed: for path in data: assert path in receiver.plots() elif expected_state == State.invalid: for path in data: assert path not in receiver.invalid() elif expected_state == State.keys_missing: for path in data: assert path not in receiver.keys_missing() elif expected_state == State.duplicates: for path in data: assert path not in receiver.duplicates() def post_function_validate(receiver: Receiver, data: Union[List[Plot], List[str]], expected_state: State) -> None: if expected_state == State.loaded: for plot_info in data: assert type(plot_info) == Plot assert plot_info.filename in receiver._current_sync.delta.valid.additions elif expected_state == State.removed: for path in data: assert path in receiver._current_sync.delta.valid.removals elif expected_state == State.invalid: for path in data: assert path in receiver._current_sync.delta.invalid.additions elif expected_state == State.keys_missing: for path in data: assert path in receiver._current_sync.delta.keys_missing.additions elif expected_state == State.duplicates: for path in data: assert path in receiver._current_sync.delta.duplicates.additions @pytest.mark.asyncio async def run_sync_step(receiver: Receiver, sync_step: SyncStepData, expected_state: State) -> None: assert receiver.current_sync().state == expected_state last_sync_time_before = receiver._last_sync.time_done # For the the list types invoke the trigger function in batches if sync_step.payload_type == PlotSyncPlotList or sync_step.payload_type == PlotSyncPathList: step_data, _ = sync_step.args assert len(step_data) == 10 # Invoke batches of: 1, 2, 3, 4 items and validate the data against plot store before and after indexes = [0, 1, 3, 6, 10] for i in range(0, len(indexes) - 1): plots_processed_before = receiver.current_sync().plots_processed invoke_data = step_data[indexes[i] : indexes[i + 1]] pre_function_validate(receiver, invoke_data, expected_state) await sync_step.function( create_payload(sync_step.payload_type, False, invoke_data, i == (len(indexes) - 2)) ) post_function_validate(receiver, invoke_data, expected_state) if expected_state == State.removed: assert receiver.current_sync().plots_processed == plots_processed_before else: assert receiver.current_sync().plots_processed == plots_processed_before + len(invoke_data) else: # For Start/Done just invoke it.. await sync_step.function(create_payload(sync_step.payload_type, sync_step.state == State.idle, *sync_step.args)) # Make sure we moved to the next state assert receiver.current_sync().state != expected_state if sync_step.payload_type == PlotSyncDone: assert receiver._last_sync.time_done != last_sync_time_before assert receiver.last_sync().plots_processed == receiver.last_sync().plots_total else: assert receiver._last_sync.time_done == last_sync_time_before def plot_sync_setup() -> Tuple[Receiver, List[SyncStepData]]: harvester_connection = get_dummy_connection(NodeType.HARVESTER) receiver = Receiver(harvester_connection, dummy_callback) # type:ignore[arg-type] # Create example plot data path_list = [str(x) for x in range(0, 40)] plot_info_list = [ Plot( filename=str(x), size=uint8(0), plot_id=bytes32(token_bytes(32)), pool_contract_puzzle_hash=None, pool_public_key=None, plot_public_key=G1Element(), file_size=uint64(random.randint(0, 100)), time_modified=uint64(0), ) for x in path_list ] # Manually add the plots we want to remove in tests receiver._plots = {plot_info.filename: plot_info for plot_info in plot_info_list[0:10]} receiver._total_plot_size = sum(plot.file_size for plot in receiver._plots.values()) sync_steps: List[SyncStepData] = [ SyncStepData(State.idle, receiver.sync_started, PlotSyncStart, False, uint64(0), uint32(len(plot_info_list))), SyncStepData(State.loaded, receiver.process_loaded, PlotSyncPlotList, plot_info_list[10:20], True), SyncStepData(State.removed, receiver.process_removed, PlotSyncPathList, path_list[0:10], True), SyncStepData(State.invalid, receiver.process_invalid, PlotSyncPathList, path_list[20:30], True), SyncStepData(State.keys_missing, receiver.process_keys_missing, PlotSyncPathList, path_list[30:40], True), SyncStepData(State.duplicates, receiver.process_duplicates, PlotSyncPathList, path_list[10:20], True), SyncStepData(State.done, receiver.sync_done, PlotSyncDone, uint64(0)), ] return receiver, sync_steps def test_default_values() -> None: assert_default_values(Receiver(get_dummy_connection(NodeType.HARVESTER), dummy_callback)) # type:ignore[arg-type] @pytest.mark.asyncio async def test_reset() -> None: receiver, sync_steps = plot_sync_setup() connection_before = receiver.connection() # Assign some dummy values receiver._current_sync.state = State.done receiver._current_sync.sync_id = uint64(1) receiver._current_sync.next_message_id = uint64(1) receiver._current_sync.plots_processed = uint32(1) receiver._current_sync.plots_total = uint32(1) receiver._current_sync.delta.valid.additions = receiver.plots().copy() receiver._current_sync.delta.valid.removals = ["1"] receiver._current_sync.delta.invalid.additions = ["1"] receiver._current_sync.delta.invalid.removals = ["1"] receiver._current_sync.delta.keys_missing.additions = ["1"] receiver._current_sync.delta.keys_missing.removals = ["1"] receiver._current_sync.delta.duplicates.additions = ["1"] receiver._current_sync.delta.duplicates.removals = ["1"] receiver._current_sync.time_done = time.time() receiver._last_sync = dataclasses.replace(receiver._current_sync) receiver._invalid = ["1"] receiver._keys_missing = ["1"] receiver._duplicates = ["1"] receiver._last_sync.sync_id = uint64(1) # Call `reset` and make sure all expected values are set back to their defaults. receiver.reset() assert_default_values(receiver) assert receiver._current_sync.delta == Delta() # Connection should remain assert receiver.connection() == connection_before @pytest.mark.parametrize("counts_only", [True, False]) @pytest.mark.asyncio async def test_to_dict(counts_only: bool) -> None: receiver, sync_steps = plot_sync_setup() plot_sync_dict_1 = receiver.to_dict(counts_only) assert get_list_or_len(plot_sync_dict_1["plots"], not counts_only) == 10 assert get_list_or_len(plot_sync_dict_1["failed_to_open_filenames"], not counts_only) == 0 assert get_list_or_len(plot_sync_dict_1["no_key_filenames"], not counts_only) == 0 assert get_list_or_len(plot_sync_dict_1["duplicates"], not counts_only) == 0 assert plot_sync_dict_1["total_plot_size"] == sum(plot.file_size for plot in receiver.plots().values()) assert plot_sync_dict_1["syncing"] is None assert plot_sync_dict_1["last_sync_time"] is None assert plot_sync_dict_1["connection"] == { "node_id": receiver.connection().peer_node_id, "host": receiver.connection().peer_host, "port": receiver.connection().peer_port, } # We should get equal dicts assert plot_sync_dict_1 == receiver.to_dict(counts_only) # But unequal dicts wit the opposite counts_only value assert plot_sync_dict_1 != receiver.to_dict(not counts_only) expected_plot_files_processed: int = 0 expected_plot_files_total: int = sync_steps[State.idle].args[2] # Walk through all states from idle to done and run them with the test data and validate the sync progress for state in State: await run_sync_step(receiver, sync_steps[state], state) if state != State.idle and state != State.removed and state != State.done: expected_plot_files_processed += len(sync_steps[state].args[0]) sync_data = receiver.to_dict()["syncing"] if state == State.done: expected_sync_data = None else: expected_sync_data = { "initial": True, "plot_files_processed": expected_plot_files_processed, "plot_files_total": expected_plot_files_total, } assert sync_data == expected_sync_data plot_sync_dict_3 = receiver.to_dict(counts_only) assert get_list_or_len(sync_steps[State.loaded].args[0], counts_only) == plot_sync_dict_3["plots"] assert ( get_list_or_len(sync_steps[State.invalid].args[0], counts_only) == plot_sync_dict_3["failed_to_open_filenames"] ) assert get_list_or_len(sync_steps[State.keys_missing].args[0], counts_only) == plot_sync_dict_3["no_key_filenames"] assert get_list_or_len(sync_steps[State.duplicates].args[0], counts_only) == plot_sync_dict_3["duplicates"] assert plot_sync_dict_3["total_plot_size"] == sum(plot.file_size for plot in receiver.plots().values()) assert plot_sync_dict_3["last_sync_time"] > 0 assert plot_sync_dict_3["syncing"] is None # Trigger a repeated plot sync await receiver.sync_started( PlotSyncStart( PlotSyncIdentifier(uint64(time.time()), uint64(receiver.last_sync().sync_id + 1), uint64(0)), False, receiver.last_sync().sync_id, uint32(1), ) ) assert receiver.to_dict()["syncing"] == { "initial": False, "plot_files_processed": 0, "plot_files_total": 1, } @pytest.mark.asyncio async def test_sync_flow() -> None: receiver, sync_steps = plot_sync_setup() for plot_info in sync_steps[State.loaded].args[0]: assert plot_info.filename not in receiver.plots() for path in sync_steps[State.removed].args[0]: assert path in receiver.plots() for path in sync_steps[State.invalid].args[0]: assert path not in receiver.invalid() for path in sync_steps[State.keys_missing].args[0]: assert path not in receiver.keys_missing() for path in sync_steps[State.duplicates].args[0]: assert path not in receiver.duplicates() # Walk through all states from idle to done and run them with the test data for state in State: await run_sync_step(receiver, sync_steps[state], state) for plot_info in sync_steps[State.loaded].args[0]: assert plot_info.filename in receiver.plots() for path in sync_steps[State.removed].args[0]: assert path not in receiver.plots() for path in sync_steps[State.invalid].args[0]: assert path in receiver.invalid() for path in sync_steps[State.keys_missing].args[0]: assert path in receiver.keys_missing() for path in sync_steps[State.duplicates].args[0]: assert path in receiver.duplicates() # We should be in idle state again assert receiver.current_sync().state == State.idle @pytest.mark.asyncio async def test_invalid_ids() -> None: receiver, sync_steps = plot_sync_setup() for state in State: assert receiver.current_sync().state == state current_step = sync_steps[state] if receiver.current_sync().state == State.idle: # Set last_sync_id for the tests below receiver._last_sync.sync_id = uint64(1) # Test "sync_started last doesn't match" invalid_last_sync_id_param = PlotSyncStart( plot_sync_identifier(uint64(0), uint64(0)), False, uint64(2), uint32(0) ) await current_step.function(invalid_last_sync_id_param) assert_error_response(receiver, ErrorCodes.invalid_last_sync_id) # Test "last_sync_id == new_sync_id" invalid_sync_id_match_param = PlotSyncStart( plot_sync_identifier(uint64(1), uint64(0)), False, uint64(1), uint32(0) ) await current_step.function(invalid_sync_id_match_param) assert_error_response(receiver, ErrorCodes.sync_ids_match) # Reset the last_sync_id to the default receiver._last_sync.sync_id = uint64(0) else: # Test invalid sync_id invalid_sync_id_param = current_step.payload_type( plot_sync_identifier(uint64(10), uint64(receiver.current_sync().next_message_id)), *current_step.args ) await current_step.function(invalid_sync_id_param) assert_error_response(receiver, ErrorCodes.invalid_identifier) # Test invalid message_id invalid_message_id_param = current_step.payload_type( plot_sync_identifier( receiver.current_sync().sync_id, uint64(receiver.current_sync().next_message_id + 1) ), *current_step.args, ) await current_step.function(invalid_message_id_param) assert_error_response(receiver, ErrorCodes.invalid_identifier) payload = create_payload(current_step.payload_type, state == State.idle, *current_step.args) await current_step.function(payload) @pytest.mark.parametrize( ["state_to_fail", "expected_error_code"], [ pytest.param(State.loaded, ErrorCodes.plot_already_available, id="already available plots"), pytest.param(State.invalid, ErrorCodes.plot_already_available, id="already available paths"), pytest.param(State.removed, ErrorCodes.plot_not_available, id="not available"), ], ) @pytest.mark.asyncio async def test_plot_errors(state_to_fail: State, expected_error_code: ErrorCodes) -> None: receiver, sync_steps = plot_sync_setup() for state in State: assert receiver.current_sync().state == state current_step = sync_steps[state] if state == state_to_fail: plot_infos, _ = current_step.args await current_step.function(create_payload(current_step.payload_type, False, plot_infos, False)) identifier = plot_sync_identifier(receiver.current_sync().sync_id, receiver.current_sync().next_message_id) invalid_payload = current_step.payload_type(identifier, plot_infos, True) await current_step.function(invalid_payload) if state == state_to_fail: assert_error_response(receiver, expected_error_code) return else: await current_step.function( create_payload(current_step.payload_type, state == State.idle, *current_step.args) ) assert False, "Didn't fail in the expected state"