from __future__ import annotations import logging import sys import time from collections.abc import Callable, Iterator from dataclasses import dataclass, replace from os import unlink from pathlib import Path from shutil import copy, move from typing import cast import pytest from chia_rs import G1Element from chia_rs.sized_ints import uint16, uint32 from chiapos import DiskProver from chia._tests.plotting.util import get_test_plots from chia._tests.util.misc import boolean_datacases from chia._tests.util.time_out_assert import time_out_assert from chia.consensus.default_constants import DEFAULT_CONSTANTS from chia.plotting.cache import CURRENT_VERSION, CacheDataV1 from chia.plotting.manager import Cache, PlotManager from chia.plotting.prover import V1Prover from chia.plotting.util import ( PlotInfo, PlotRefreshEvents, PlotRefreshResult, add_plot_directory, get_plot_directories, remove_plot, remove_plot_directory, ) from chia.simulator.block_tools import get_plot_dir from chia.util.config import create_default_chia_config, lock_and_load_config, save_config from chia.util.streamable import VersionedBlob log = logging.getLogger(__name__) # These tests are not dependent on consensus rules pytestmark = [pytest.mark.limit_consensus_modes(reason="does not depend on consensus rules")] @dataclass class MockDiskProver: filename: str def get_filename(self) -> str: return self.filename @dataclass class MockPlotInfo: prover: MockDiskProver class Directory: path: Path plots: list[Path] def __init__(self, path: Path, plots_origin: list[Path]): self.path = path path.mkdir(parents=True, exist_ok=True) # Drop the existing files in the test directories for plot in path.iterdir(): unlink(plot) # Copy over the original plots for plot in plots_origin: if not Path(path / plot.name).exists(): copy(plot, path) # Adjust the paths to reflect the testing plots self.plots = [path / plot.name for plot in plots_origin] def __len__(self): return len(self.plots) def plot_info_list(self) -> list[MockPlotInfo]: return [MockPlotInfo(MockDiskProver(str(x))) for x in self.plots] def path_list(self) -> list[Path]: return self.plots def drop(self, path: Path): assert self.path / path.name del self.plots[self.plots.index(self.path / path.name)] class PlotRefreshTester: plot_manager: PlotManager expected_result: PlotRefreshResult expected_result_matched: bool def __init__(self, root_path: Path): self.plot_manager = PlotManager(root_path, self.refresh_callback, DEFAULT_CONSTANTS) # Set a very high refresh interval here to avoid unintentional refresh cycles self.plot_manager.refresh_parameter = replace( self.plot_manager.refresh_parameter, interval_seconds=uint32(10000) ) # Set to the current time to avoid automated refresh after we start below. self.plot_manager.last_refresh_time = time.time() self.plot_manager.start_refreshing() def refresh_callback(self, event: PlotRefreshEvents, refresh_result: PlotRefreshResult): if event != PlotRefreshEvents.done: # Only validate the final results for this tests return for name in ["loaded", "removed", "processed", "remaining"]: try: actual_value = refresh_result.__getattribute__(name) if type(actual_value) is list: expected_list = self.expected_result.__getattribute__(name) if len(expected_list) != len(actual_value): return values_found = 0 for value in actual_value: if type(value) is PlotInfo: for plot_info in expected_list: if plot_info.prover.get_filename() == value.prover.get_filename(): values_found += 1 continue elif value in expected_list: values_found += 1 continue if values_found != len(expected_list): log.error(f"{name} invalid: values_found {values_found} expected {len(expected_list)}") return else: expected_value = self.expected_result.__getattribute__(name) if actual_value != expected_value: log.error(f"{name} invalid: actual {actual_value} expected {expected_value}") return except AttributeError as error: log.error(f"{error}") return self.expected_result_matched = True async def run(self, expected_result: PlotRefreshResult): self.expected_result = expected_result self.expected_result_matched = False self.plot_manager.trigger_refresh() await time_out_assert(5, self.plot_manager.needs_refresh, value=False) assert self.expected_result_matched @dataclass class Environment: root_path: Path refresh_tester: PlotRefreshTester dir_1: Directory dir_2: Directory @pytest.fixture(scope="function") def environment(tmp_path, bt) -> Iterator[Environment]: dir_1_count: int = 7 dir_2_count: int = 3 plots: list[Path] = get_test_plots() assert len(plots) >= dir_1_count + dir_2_count dir_1: Directory = Directory(tmp_path / "plots" / "1", plots[0:dir_1_count]) dir_2: Directory = Directory(tmp_path / "plots" / "2", plots[dir_1_count : dir_1_count + dir_2_count]) create_default_chia_config(tmp_path) refresh_tester = PlotRefreshTester(tmp_path) refresh_tester.plot_manager.set_public_keys(bt.plot_manager.farmer_public_keys, bt.plot_manager.pool_public_keys) yield Environment(tmp_path, refresh_tester, dir_1, dir_2) refresh_tester.plot_manager.stop_refreshing() # Wrap `remove_plot` to give it the same interface as the other triggers, e.g. `add_plot_directory(Path, str)`. def trigger_remove_plot(_: Path, plot_path: str): remove_plot(Path(plot_path)) @pytest.mark.limit_consensus_modes(reason="not dependent on consensus, does not support parallel execution") @pytest.mark.anyio async def test_plot_refreshing(environment): env: Environment = environment expected_result = PlotRefreshResult() dir_duplicates: Directory = Directory(get_plot_dir().resolve() / "duplicates", env.dir_1.plots) async def run_test_case( *, trigger: Callable, test_path: Path, expect_loaded: list[MockPlotInfo], expect_removed: list[Path], expect_processed: int, expect_duplicates: int, expected_directories: int, expect_total_plots: int, ): expected_result.loaded = cast(list[PlotInfo], expect_loaded) expected_result.removed = expect_removed expected_result.processed = expect_processed trigger(env.root_path, str(test_path)) assert len(get_plot_directories(env.root_path)) == expected_directories await env.refresh_tester.run(expected_result) assert len(env.refresh_tester.plot_manager.plots) == expect_total_plots assert len(env.refresh_tester.plot_manager.get_duplicates()) == expect_duplicates assert len(env.refresh_tester.plot_manager.failed_to_open_filenames) == 0 # Add dir_1 await run_test_case( trigger=add_plot_directory, test_path=env.dir_1.path, expect_loaded=env.dir_1.plot_info_list(), expect_removed=[], expect_processed=len(env.dir_1), expect_duplicates=0, expected_directories=1, expect_total_plots=len(env.dir_1), ) # Add dir_2 await run_test_case( trigger=add_plot_directory, test_path=env.dir_2.path, expect_loaded=env.dir_2.plot_info_list(), expect_removed=[], expect_processed=len(env.dir_1) + len(env.dir_2), expect_duplicates=0, expected_directories=2, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) # Add dir_duplicates await run_test_case( trigger=add_plot_directory, test_path=dir_duplicates.path, expect_loaded=[], expect_removed=[], expect_processed=len(env.dir_1) + len(env.dir_2) + len(dir_duplicates), expect_duplicates=len(dir_duplicates), expected_directories=3, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) for item in dir_duplicates.path.iterdir(): assert item.is_file() and item in env.refresh_tester.plot_manager.get_duplicates() # Drop the duplicated plot we remove in the next test case from the test directory upfront so that the numbers match # the expected below drop_path = dir_duplicates.plots[0] dir_duplicates.drop(drop_path) # Delete one duplicated plot await run_test_case( trigger=trigger_remove_plot, test_path=drop_path, expect_loaded=[], expect_removed=[], expect_processed=len(env.dir_1) + len(env.dir_2) + len(dir_duplicates), expect_duplicates=len(dir_duplicates), expected_directories=3, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) # Drop the duplicated plot we remove in the next test case from the test directory upfront so that the numbers match # the expected below drop_path = env.dir_1.plots[0] env.dir_1.drop(drop_path) # Delete one duplicated plot await run_test_case( trigger=trigger_remove_plot, test_path=drop_path, expect_loaded=[], expect_removed=[drop_path], expect_processed=len(env.dir_1) + len(env.dir_2) + len(dir_duplicates), expect_duplicates=len(dir_duplicates), expected_directories=3, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) # Remove directory with the duplicates await run_test_case( trigger=remove_plot_directory, test_path=dir_duplicates.path, expect_loaded=[], expect_removed=[], expect_processed=len(env.dir_1) + len(env.dir_2), expect_duplicates=0, expected_directories=2, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) for item in dir_duplicates.path.iterdir(): assert item.is_file() and item not in env.refresh_tester.plot_manager.get_duplicates() # Re-add the directory with the duplicates for other tests await run_test_case( trigger=add_plot_directory, test_path=dir_duplicates.path, expect_loaded=[], expect_removed=[], expect_processed=len(env.dir_1) + len(env.dir_2) + len(dir_duplicates), expect_duplicates=len(dir_duplicates), expected_directories=3, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) # Remove dir_1 from which the duplicated plots are loaded. This removes the duplicates of dir_1 # and in the same run loads them from dir_duplicates. await run_test_case( trigger=remove_plot_directory, test_path=env.dir_1.path, expect_loaded=dir_duplicates.plot_info_list(), expect_removed=env.dir_1.path_list(), expect_processed=len(env.dir_2) + len(dir_duplicates), expect_duplicates=0, expected_directories=2, expect_total_plots=len(env.dir_2) + len(dir_duplicates), ) # Re-add the directory. Now the plot seen as duplicate is from dir_1, not from dir_duplicates like before await run_test_case( trigger=add_plot_directory, test_path=env.dir_1.path, expect_loaded=[], expect_removed=[], expect_processed=len(env.dir_1) + len(env.dir_2) + len(dir_duplicates), expect_duplicates=len(dir_duplicates), expected_directories=3, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) # Drop the duplicated plot we remove in the next test case from the test directory upfront so that the numbers match # the expected below drop_path = env.dir_1.plots[2] env.dir_1.drop(drop_path) # Remove the duplicated plot await run_test_case( trigger=trigger_remove_plot, test_path=drop_path, expect_loaded=[], expect_removed=[], expect_processed=len(env.dir_1) + len(env.dir_2) + len(dir_duplicates), expect_duplicates=len(env.dir_1), expected_directories=3, expect_total_plots=len(env.dir_2) + len(dir_duplicates), ) # Remove dir_duplicates, this drops the duplicates and loads all plots from dir_1 await run_test_case( trigger=remove_plot_directory, test_path=dir_duplicates.path, expect_loaded=env.dir_1.plot_info_list(), expect_removed=dir_duplicates.path_list(), expect_processed=len(env.dir_1) + len(env.dir_2), expect_duplicates=0, expected_directories=2, expect_total_plots=len(env.dir_1) + len(env.dir_2), ) # Remove dir_2 await run_test_case( trigger=remove_plot_directory, test_path=env.dir_2.path, expect_loaded=[], expect_removed=env.dir_2.path_list(), expect_processed=len(env.dir_1), expect_duplicates=0, expected_directories=1, expect_total_plots=len(env.dir_1), ) # Remove dir_1 await run_test_case( trigger=remove_plot_directory, test_path=env.dir_1.path, expect_loaded=[], expect_removed=env.dir_1.path_list(), expect_processed=0, expect_duplicates=0, expected_directories=0, expect_total_plots=0, ) @pytest.mark.anyio async def test_initial_refresh_flag(environment: Environment) -> None: env: Environment = environment assert env.refresh_tester.plot_manager.initial_refresh() for _ in range(2): await env.refresh_tester.run(PlotRefreshResult()) assert not env.refresh_tester.plot_manager.initial_refresh() env.refresh_tester.plot_manager.reset() assert env.refresh_tester.plot_manager.initial_refresh() @pytest.mark.anyio async def test_invalid_plots(environment): env: Environment = environment expected_result = PlotRefreshResult() # Test re-trying if processing a plot failed # First create a backup of the plot retry_test_plot = env.dir_1.path_list()[0].resolve() retry_test_plot_save = Path(env.dir_1.path / ".backup").resolve() copy(retry_test_plot, retry_test_plot_save) # Invalidate the plot with open(retry_test_plot, "r+b") as file: file.write(bytes(100)) # Add it and validate it fails to load add_plot_directory(env.root_path, str(env.dir_1.path)) expected_result.loaded = env.dir_1.plot_info_list()[1:] expected_result.removed = [] expected_result.processed = len(env.dir_1) expected_result.remaining = 0 await env.refresh_tester.run(expected_result) assert len(env.refresh_tester.plot_manager.failed_to_open_filenames) == 1 assert retry_test_plot in env.refresh_tester.plot_manager.failed_to_open_filenames # Give it a non .plot ending and make sure it gets removed from the invalid list on the next refresh retry_test_plot_unload = Path(env.dir_1.path / ".unload").resolve() move(retry_test_plot, retry_test_plot_unload) expected_result.processed -= 1 expected_result.loaded = [] await env.refresh_tester.run(expected_result) assert len(env.refresh_tester.plot_manager.failed_to_open_filenames) == 0 assert retry_test_plot not in env.refresh_tester.plot_manager.failed_to_open_filenames # Recover the name and make sure it reappears in the invalid list move(retry_test_plot_unload, retry_test_plot) expected_result.processed += 1 await env.refresh_tester.run(expected_result) assert len(env.refresh_tester.plot_manager.failed_to_open_filenames) == 1 assert retry_test_plot in env.refresh_tester.plot_manager.failed_to_open_filenames # Make sure the file stays in `failed_to_open_filenames` and doesn't get loaded in the next refresh cycle expected_result.loaded = [] expected_result.processed = len(env.dir_1) await env.refresh_tester.run(expected_result) assert len(env.refresh_tester.plot_manager.failed_to_open_filenames) == 1 assert retry_test_plot in env.refresh_tester.plot_manager.failed_to_open_filenames # Now decrease the re-try timeout, restore the valid plot file and make sure it properly loads now env.refresh_tester.plot_manager.refresh_parameter = replace( env.refresh_tester.plot_manager.refresh_parameter, retry_invalid_seconds=uint32(0) ) move(retry_test_plot_save, retry_test_plot) expected_result.loaded = env.dir_1.plot_info_list()[0:1] expected_result.processed = len(env.dir_1) await env.refresh_tester.run(expected_result) assert len(env.refresh_tester.plot_manager.failed_to_open_filenames) == 0 assert retry_test_plot not in env.refresh_tester.plot_manager.failed_to_open_filenames @pytest.mark.anyio async def test_keys_missing(environment: Environment) -> None: env: Environment = environment not_in_keychain_plots: list[Path] = get_test_plots("not_in_keychain") dir_not_in_keychain: Directory = Directory(env.root_path / "plots" / "not_in_keychain", not_in_keychain_plots) expected_result = PlotRefreshResult() # The plots in "not_in_keychain" directory have infinity g1 elements as farmer/pool key so they should be plots # with missing keys for now add_plot_directory(env.root_path, str(dir_not_in_keychain.path)) expected_result.loaded = [] expected_result.removed = [] expected_result.processed = len(dir_not_in_keychain) expected_result.remaining = 0 for i in range(2): await env.refresh_tester.run(expected_result) assert len(env.refresh_tester.plot_manager.no_key_filenames) == len(dir_not_in_keychain) for path in env.refresh_tester.plot_manager.no_key_filenames: assert path in dir_not_in_keychain.plots # Delete one of the plots and make sure it gets dropped from the no key filenames list drop_plot = dir_not_in_keychain.path_list()[0] dir_not_in_keychain.drop(drop_plot) drop_plot.unlink() assert drop_plot in env.refresh_tester.plot_manager.no_key_filenames expected_result.processed -= 1 await env.refresh_tester.run(expected_result) assert drop_plot not in env.refresh_tester.plot_manager.no_key_filenames # Now add the missing keys to the plot manager's key lists and make sure the plots are getting loaded env.refresh_tester.plot_manager.farmer_public_keys.append(G1Element()) env.refresh_tester.plot_manager.pool_public_keys.append(G1Element()) expected_result.loaded = dir_not_in_keychain.plot_info_list() # type: ignore[assignment] expected_result.processed = len(dir_not_in_keychain) await env.refresh_tester.run(expected_result) # And make sure they are dropped from the list of plots with missing keys assert len(env.refresh_tester.plot_manager.no_key_filenames) == 0 @pytest.mark.anyio async def test_plot_info_caching(environment, bt): env: Environment = environment expected_result = PlotRefreshResult() add_plot_directory(env.root_path, str(env.dir_1.path)) expected_result.loaded = env.dir_1.plot_info_list() expected_result.removed = [] expected_result.processed = len(env.dir_1) expected_result.remaining = 0 await env.refresh_tester.run(expected_result) assert env.refresh_tester.plot_manager.cache.path().exists() unlink(env.refresh_tester.plot_manager.cache.path()) # Should not write the cache again on shutdown because it didn't change assert not env.refresh_tester.plot_manager.cache.path().exists() env.refresh_tester.plot_manager.stop_refreshing() assert not env.refresh_tester.plot_manager.cache.path().exists() # Manually trigger `save_cache` and make sure it creates a new cache file env.refresh_tester.plot_manager.cache.save() assert env.refresh_tester.plot_manager.cache.path().exists() refresh_tester: PlotRefreshTester = PlotRefreshTester(env.root_path) plot_manager = refresh_tester.plot_manager plot_manager.set_public_keys(bt.plot_manager.farmer_public_keys, bt.plot_manager.pool_public_keys) plot_manager.cache.load() assert len(plot_manager.cache) == len(env.refresh_tester.plot_manager.cache) for path, cache_entry in env.refresh_tester.plot_manager.cache.items(): cache_entry_new = plot_manager.cache.get(path) assert bytes(cache_entry_new.prover) == bytes(cache_entry.prover) assert cache_entry_new.farmer_public_key == cache_entry.farmer_public_key assert cache_entry_new.pool_public_key == cache_entry.pool_public_key assert cache_entry_new.pool_contract_puzzle_hash == cache_entry.pool_contract_puzzle_hash assert cache_entry_new.plot_public_key == cache_entry.plot_public_key await refresh_tester.run(expected_result) for path, plot_info in env.refresh_tester.plot_manager.plots.items(): assert path in plot_manager.plots assert plot_manager.plots[path].prover.get_filename() == plot_info.prover.get_filename() assert plot_manager.plots[path].prover.get_id() == plot_info.prover.get_id() assert plot_manager.plots[path].prover.get_memo() == plot_info.prover.get_memo() assert plot_manager.plots[path].prover.get_param().size_v1 == plot_info.prover.get_param().size_v1 assert plot_manager.plots[path].prover.get_param().strength_v2 == plot_info.prover.get_param().strength_v2 assert plot_manager.plots[path].prover.get_compression_level() == plot_info.prover.get_compression_level() assert plot_manager.plots[path].pool_public_key == plot_info.pool_public_key assert plot_manager.plots[path].pool_contract_puzzle_hash == plot_info.pool_contract_puzzle_hash assert plot_manager.plots[path].plot_public_key == plot_info.plot_public_key assert plot_manager.plots[path].file_size == plot_info.file_size assert plot_manager.plots[path].time_modified == plot_info.time_modified assert plot_manager.plot_filename_paths == env.refresh_tester.plot_manager.plot_filename_paths assert plot_manager.failed_to_open_filenames == env.refresh_tester.plot_manager.failed_to_open_filenames assert plot_manager.no_key_filenames == env.refresh_tester.plot_manager.no_key_filenames plot_manager.stop_refreshing() # Modify the content of the plot_manager.dat with open(plot_manager.cache.path(), "r+b") as file: file.write(b"\xff\xff") # Sets Cache.version to 65535 # Make sure it just loads the plots normally if it fails to load the cache refresh_tester: PlotRefreshTester = PlotRefreshTester(env.root_path) plot_manager = refresh_tester.plot_manager plot_manager.cache.load() assert len(plot_manager.cache) == 0 plot_manager.set_public_keys(bt.plot_manager.farmer_public_keys, bt.plot_manager.pool_public_keys) await refresh_tester.run(expected_result) assert len(plot_manager.plots) == len(plot_manager.plots) plot_manager.stop_refreshing() @pytest.mark.anyio async def test_drop_too_large_cache_entries(environment, bt): env: Environment = environment expected_result = PlotRefreshResult(loaded=env.dir_1.plot_info_list(), processed=len(env.dir_1)) add_plot_directory(env.root_path, str(env.dir_1.path)) await env.refresh_tester.run(expected_result) assert env.refresh_tester.plot_manager.cache.path().exists() assert len(env.dir_1) >= 6, "This test requires at least 6 cache entries" # Load the cache entries cache_path = env.refresh_tester.plot_manager.cache.path() serialized = cache_path.read_bytes() stored_cache: VersionedBlob = VersionedBlob.from_bytes(serialized) cache_data: CacheDataV1 = CacheDataV1.from_bytes(stored_cache.blob) def modify_cache_entry(index: int, additional_data: int, modify_memo: bool) -> str: path, cache_entry = cache_data.entries[index] prover_data = cache_entry.prover_data # Size of length hints in chiapos serialization currently depends on the platform size_length = 8 if sys.maxsize > 2**32 else 4 # Version version_size = 2 version = prover_data[0:version_size] # Filename filename_offset = version_size + size_length filename_length = int.from_bytes(prover_data[version_size:filename_offset], byteorder=sys.byteorder) filename = prover_data[filename_offset : filename_offset + filename_length] # Memo memo_length_offset = filename_offset + filename_length memo_length = int.from_bytes( prover_data[memo_length_offset : memo_length_offset + size_length], byteorder=sys.byteorder ) memo_offset = memo_length_offset + size_length memo = prover_data[memo_offset : memo_offset + memo_length] # id, k, table pointers, C2 remainder = prover_data[memo_offset + memo_length :] # Add the additional data to the filename filename_length += additional_data filename += bytes(b"\a" * additional_data) # Add the additional data to the memo if requested if modify_memo: memo_length += additional_data memo += bytes(b"\b" * additional_data) filename_length_bytes = filename_length.to_bytes(size_length, byteorder=sys.byteorder) memo_length_bytes = memo_length.to_bytes(size_length, byteorder=sys.byteorder) cache_data.entries[index] = ( path, replace( cache_entry, prover_data=bytes(version + filename_length_bytes + filename + memo_length_bytes + memo + remainder), ), ) return path def assert_cache(expected: list[MockPlotInfo]) -> None: test_cache = Cache(cache_path) assert len(test_cache) == 0 test_cache.load() assert len(test_cache) == len(expected) for plot_info in expected: assert test_cache.get(Path(plot_info.prover.get_filename())) is not None # Modify two entries, with and without memo modification, they both should remain in the cache after load modify_cache_entry(0, 1500, modify_memo=False) modify_cache_entry(1, 1500, modify_memo=True) invalid_entries = [ modify_cache_entry(2, 2000, modify_memo=False), modify_cache_entry(3, 2000, modify_memo=True), modify_cache_entry(4, 50000, modify_memo=False), modify_cache_entry(5, 50000, modify_memo=True), ] plot_infos = env.dir_1.plot_info_list() # Make sure the cache currently contains all plots from dir1 assert_cache(plot_infos) # Write the modified cache entries to the file cache_path.write_bytes(bytes(VersionedBlob(uint16(CURRENT_VERSION), bytes(cache_data)))) # And now test that plots in invalid_entries are not longer loaded assert_cache([plot_info for plot_info in plot_infos if plot_info.prover.get_filename() not in invalid_entries]) @pytest.mark.anyio async def test_cache_lifetime(environment: Environment) -> None: # Load a directory to produce a cache file env: Environment = environment expected_result = PlotRefreshResult() add_plot_directory(env.root_path, str(env.dir_1.path)) expected_result.loaded = env.dir_1.plot_info_list() # type: ignore[assignment] expected_result.removed = [] expected_result.processed = len(env.dir_1) expected_result.remaining = 0 await env.refresh_tester.run(expected_result) expected_result.loaded = [] cache_v1: Cache = env.refresh_tester.plot_manager.cache assert len(cache_v1) > 0 count_before = len(cache_v1) # Remove half of the plots in dir1 for path in env.dir_1.path_list()[0 : int(len(env.dir_1) / 2)]: expected_result.processed -= 1 expected_result.removed.append(path) unlink(path) # Modify the `last_use` timestamp of all cache entries to let them expire last_use_before = time.time() - Cache.expiry_seconds - 1 for cache_entry in cache_v1.values(): cache_entry.last_use = last_use_before assert cache_entry.expired(Cache.expiry_seconds) # The next refresh cycle will now lead to half of the cache entries being removed because they are expired and # the related plots do not longer exist. await env.refresh_tester.run(expected_result) assert len(cache_v1) == count_before - len(expected_result.removed) # The other half of the cache entries should have a different `last_use` value now. for cache_entry in cache_v1.values(): assert cache_entry.last_use != last_use_before @pytest.mark.parametrize( ["event_to_raise"], [ pytest.param(PlotRefreshEvents.started, id="started"), pytest.param(PlotRefreshEvents.batch_processed, id="batch_processed"), pytest.param(PlotRefreshEvents.done, id="done"), ], ) @pytest.mark.anyio async def test_callback_event_raises(environment, event_to_raise: PlotRefreshEvents): last_event_fired: PlotRefreshEvents | None = None def raising_callback(event: PlotRefreshEvents, _: PlotRefreshResult): nonlocal last_event_fired last_event_fired = event if event == event_to_raise: raise Exception(f"run_raise_in_callback {event_to_raise}") env: Environment = environment expected_result = PlotRefreshResult() # Load dir_1 add_plot_directory(env.root_path, str(env.dir_1.path)) expected_result.loaded = env.dir_1.plot_info_list() # type: ignore[assignment] expected_result.removed = [] expected_result.processed = len(env.dir_1) expected_result.remaining = 0 await env.refresh_tester.run(expected_result) # Load dir_2 add_plot_directory(env.root_path, str(env.dir_2.path)) expected_result.loaded = env.dir_2.plot_info_list() # type: ignore[assignment] expected_result.removed = [] expected_result.processed = len(env.dir_1) + len(env.dir_2) expected_result.remaining = 0 await env.refresh_tester.run(expected_result) # Now raise the exception in the callback default_callback = env.refresh_tester.plot_manager._refresh_callback env.refresh_tester.plot_manager.set_refresh_callback(raising_callback) env.refresh_tester.plot_manager.start_refreshing() env.refresh_tester.plot_manager.trigger_refresh() await time_out_assert(5, env.refresh_tester.plot_manager.needs_refresh, value=False) # And make sure the follow-up evens aren't fired assert last_event_fired == event_to_raise # The exception should trigger `PlotManager.reset()` and clear the plots assert len(env.refresh_tester.plot_manager.plots) == 0 assert len(env.refresh_tester.plot_manager.plot_filename_paths) == 0 assert len(env.refresh_tester.plot_manager.failed_to_open_filenames) == 0 assert len(env.refresh_tester.plot_manager.no_key_filenames) == 0 # The next run without the valid callback should lead to re-loading of all plot env.refresh_tester.plot_manager.set_refresh_callback(default_callback) expected_result.loaded = env.dir_1.plot_info_list() + env.dir_2.plot_info_list() # type: ignore[assignment] expected_result.removed = [] expected_result.processed = len(env.dir_1) + len(env.dir_2) expected_result.remaining = 0 await env.refresh_tester.run(expected_result) @pytest.mark.anyio async def test_recursive_plot_scan(environment: Environment) -> None: env: Environment = environment # Create a directory tree with some subdirectories containing plots, others not. root_plot_dir = env.root_path / "root" sub_dir_0: Directory = Directory(root_plot_dir / "0", env.dir_1.plots[0:2]) sub_dir_0_1: Directory = Directory(sub_dir_0.path / "1", env.dir_1.plots[2:3]) sub_dir_1: Directory = Directory(root_plot_dir / "1", []) sub_dir_1_0: Directory = Directory(sub_dir_1.path / "0", []) sub_dir_1_0_1: Directory = Directory(sub_dir_1_0.path / "1", env.dir_1.plots[3:7]) # List of all the plots in the directory tree expected_plot_list = sub_dir_0.plot_info_list() + sub_dir_0_1.plot_info_list() + sub_dir_1_0_1.plot_info_list() # Adding the root without `recursive_plot_scan` and running a test should not load any plots (match an empty result) expected_result = PlotRefreshResult() add_plot_directory(env.root_path, str(root_plot_dir)) await env.refresh_tester.run(expected_result) # Set the recursive scan flag in the config with lock_and_load_config(env.root_path, "config.yaml") as config: config["harvester"]["recursive_plot_scan"] = True save_config(env.root_path, "config.yaml", config) # With the flag enabled it should load all expected plots expected_result.loaded = expected_plot_list # type: ignore[assignment] expected_result.processed = len(expected_plot_list) await env.refresh_tester.run(expected_result) # Adding the subdirectories also should not lead to some failure or duplicated loading add_plot_directory(env.root_path, str(sub_dir_0_1.path)) add_plot_directory(env.root_path, str(sub_dir_1_0_1.path)) expected_result.loaded = [] await env.refresh_tester.run(expected_result) @pytest.mark.anyio async def test_disk_prover_from_bytes(environment: Environment): env: Environment = environment expected_result = PlotRefreshResult() expected_result.loaded = env.dir_1.plot_info_list() # type: ignore[assignment] expected_result.processed = len(env.dir_1) add_plot_directory(env.root_path, str(env.dir_1.path)) await env.refresh_tester.run(expected_result) _, plot_info = next(iter(env.refresh_tester.plot_manager.plots.items())) recreated_prover = V1Prover(DiskProver.from_bytes(bytes(plot_info.prover))) assert recreated_prover.get_id() == plot_info.prover.get_id() assert recreated_prover.get_filename() == plot_info.prover.get_filename() @boolean_datacases(name="follow_links", false="no_follow", true="follow") @pytest.mark.anyio async def test_recursive_plot_scan_symlinks(environment: Environment, follow_links: bool) -> None: env: Environment = environment # Create a directory tree with some subdirectories containing plots, others not. root_plot_dir = env.root_path / "root" sub_dir_0: Directory = Directory(root_plot_dir / "0", env.dir_1.plots[0:2]) sub_dir_0_1: Directory = Directory(sub_dir_0.path / "1", env.dir_1.plots[2:3]) # Create a plot directory outside of the root plot directory other_plot_dir: Directory = Directory(env.root_path / "other", env.dir_1.plots[3:7]) # Create a symlink to the other plot directory from inside the root plot directory symlink_path = Path(root_plot_dir / "other") symlink_path.symlink_to(env.root_path / "other") # Adding the root without `recursive_plot_scan` and running a test should not load any plots (match an empty result) expected_result = PlotRefreshResult() add_plot_directory(env.root_path, str(root_plot_dir)) await env.refresh_tester.run(expected_result) if follow_links: expected_plot_list = sub_dir_0.plot_info_list() + sub_dir_0_1.plot_info_list() + other_plot_dir.plot_info_list() else: expected_plot_list = sub_dir_0.plot_info_list() + sub_dir_0_1.plot_info_list() # Set the recursive scan flag in the config and symlink flag with lock_and_load_config(env.root_path, "config.yaml") as config: config["harvester"]["recursive_plot_scan"] = True config["harvester"]["recursive_follow_links"] = follow_links save_config(env.root_path, "config.yaml", config) # With the flag enabled it should load all expected plots expected_result.loaded = expected_plot_list # type: ignore[assignment] expected_result.processed = len(expected_plot_list) await env.refresh_tester.run(expected_result)