PoS2-prover (#20159)

* bump chia_rs to 0.31.0

* integrate new pos2 API from chia_rs into V2Prover

* fixup test_harvester_api mocking of Prover

* review-comments: rename get_quality() to get_string() on Quality protocol

* reivew comments

* fixup bump chia_rs

* restore quality string to bytes32

Author: arvidn

chia/_tests/core/custom_types/test_proof_of_space.py

@@ -13,7 +13,6 @@
 from chia.consensus.default_constants import DEFAULT_CONSTANTS
 from chia.types.blockchain_format.proof_of_space import (
     calculate_prefix_bits,
-    calculate_required_plot_strength,
     check_plot_size,
     make_pos,
     passes_plot_filter,
@@ -200,29 +199,6 @@ def test_verify_and_get_quality_string_v2(caplog: pytest.LogCaptureFixture, case
         assert len(caplog.text) == 0 if case.expected_error is None else case.expected_error in caplog.text
 
 
-@pytest.mark.parametrize(
-    "height, strength",
-    [
-        (0, 2),
-        (DEFAULT_CONSTANTS.HARD_FORK_HEIGHT, 2),
-        (DEFAULT_CONSTANTS.HARD_FORK2_HEIGHT, 2),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_4_HEIGHT - 1, 2),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_4_HEIGHT, 4),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_5_HEIGHT - 1, 4),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_5_HEIGHT, 5),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_6_HEIGHT - 1, 5),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_6_HEIGHT, 6),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_7_HEIGHT - 1, 6),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_7_HEIGHT, 7),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_8_HEIGHT - 1, 7),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_8_HEIGHT, 8),
-        (DEFAULT_CONSTANTS.PLOT_STRENGTH_8_HEIGHT + 1000000, 8),
-    ],
-)
-def test_calculate_plot_strength(height: uint32, strength: uint8) -> None:
-    assert calculate_required_plot_strength(DEFAULT_CONSTANTS, height) == strength
-
-
 @pytest.mark.parametrize(
     "size, valid",
     [
@@ -272,12 +248,9 @@ def test_can_create_proof(self, prefix_bits: int, seeded_random: random.Random)
 
 
 @pytest.mark.parametrize("height,expected", [(0, 3), (5496000, 2), (10542000, 1), (15592000, 0), (20643000, 0)])
-@pytest.mark.parametrize("plot_size", [PlotSize.make_v1(32), PlotSize.make_v2(28)])
-def test_calculate_prefix_bits_clamp_zero(height: uint32, expected: int, plot_size: PlotSize) -> None:
+def test_calculate_prefix_bits_clamp_zero_v1(height: uint32, expected: int) -> None:
     constants = DEFAULT_CONSTANTS.replace(NUMBER_ZERO_BITS_PLOT_FILTER_V1=uint8(3))
-    if plot_size.size_v2 is not None:
-        expected = constants.NUMBER_ZERO_BITS_PLOT_FILTER_V2
-    assert calculate_prefix_bits(constants, height, plot_size) == expected
+    assert calculate_prefix_bits(constants, height, PlotSize.make_v1(32)) == expected
 
 
 @pytest.mark.parametrize(
@@ -294,9 +267,20 @@ def test_calculate_prefix_bits_clamp_zero(height: uint32, expected: int, plot_si
         (20643000, 5),
     ],
 )
-@pytest.mark.parametrize("plot_size", [PlotSize.make_v1(32), PlotSize.make_v2(28)])
-def test_calculate_prefix_bits_default(height: uint32, expected: int, plot_size: PlotSize) -> None:
-    constants = DEFAULT_CONSTANTS
-    if plot_size.size_v2 is not None:
-        expected = DEFAULT_CONSTANTS.NUMBER_ZERO_BITS_PLOT_FILTER_V2
-    assert calculate_prefix_bits(constants, height, plot_size) == expected
+def test_calculate_prefix_bits_v1(height: uint32, expected: int) -> None:
+    assert calculate_prefix_bits(DEFAULT_CONSTANTS, height, PlotSize.make_v1(32)) == expected
+
+
+@pytest.mark.parametrize(
+    argnames=["height", "expected"],
+    argvalues=[
+        (0, 5),
+        (0xFFFFFFFA, 5),
+        (0xFFFFFFFB, 6),
+        (0xFFFFFFFC, 7),
+        (0xFFFFFFFD, 8),
+        (0xFFFFFFFF, 8),
+    ],
+)
+def test_calculate_prefix_bits_v2(height: uint32, expected: int) -> None:
+    assert calculate_prefix_bits(DEFAULT_CONSTANTS, height, PlotSize.make_v2(28)) == expected

chia/_tests/farmer_harvester/test_farmer_harvester.py

@@ -338,9 +338,11 @@ async def test_v2_partial_proofs_new_sp_hash(
         challenge_hash=bytes32(b"2" * 32),
         sp_hash=sp_hash,
         plot_identifier="test_plot_id",
-        partial_proofs=[b"test_partial_proof_1"],
+        partial_proofs=[[uint64(1), uint64(2), uint64(3), uint64(4)]],
         signage_point_index=uint8(0),
         plot_size=uint8(32),
+        strength=uint8(5),
+        plot_id=bytes32.fromhex("abababababababababababababababababababababababababababababababab"),
         pool_public_key=None,
         pool_contract_puzzle_hash=bytes32(b"4" * 32),
         plot_public_key=G1Element(),
@@ -367,9 +369,11 @@ async def test_v2_partial_proofs_missing_sp_hash(
         challenge_hash=bytes32(b"2" * 32),
         sp_hash=sp_hash,
         plot_identifier="test_plot_id",
-        partial_proofs=[b"test_partial_proof_1"],
+        partial_proofs=[[uint64(1), uint64(2), uint64(3), uint64(4)]],
         signage_point_index=uint8(0),
         plot_size=uint8(32),
+        plot_id=bytes32.fromhex("abababababababababababababababababababababababababababababababab"),
+        strength=uint8(5),
         pool_public_key=None,
         pool_contract_puzzle_hash=bytes32(b"4" * 32),
         plot_public_key=G1Element(),
@@ -409,9 +413,11 @@ async def test_v2_partial_proofs_with_existing_sp(
         challenge_hash=challenge_hash,
         sp_hash=sp_hash,
         plot_identifier="test_plot_id",
-        partial_proofs=[b"test_partial_proof_1", b"test_partial_proof_2"],
+        partial_proofs=[[uint64(1), uint64(2), uint64(3), uint64(4)], [uint64(2), uint64(3), uint64(4), uint64(5)]],
         signage_point_index=uint8(0),
         plot_size=uint8(32),
+        plot_id=bytes32.fromhex("abababababababababababababababababababababababababababababababab"),
+        strength=uint8(5),
         pool_public_key=G1Element(),
         pool_contract_puzzle_hash=bytes32(b"4" * 32),
         plot_public_key=G1Element(),
@@ -441,9 +447,11 @@ async def test_solution_response_handler(
         challenge_hash=challenge_hash,
         sp_hash=sp_hash,
         plot_identifier="test_plot_id",
-        partial_proofs=[b"test_partial_proof_for_quality"],
+        partial_proofs=[[uint64(1), uint64(2), uint64(3), uint64(4)]],
         signage_point_index=uint8(0),
         plot_size=uint8(32),
+        plot_id=bytes32.fromhex("abababababababababababababababababababababababababababababababab"),
+        strength=uint8(5),
         pool_public_key=G1Element(),
         pool_contract_puzzle_hash=bytes32(b"4" * 32),
         plot_public_key=G1Element(),
@@ -452,7 +460,8 @@ async def test_solution_response_handler(
     harvester_peer = await get_harvester_peer(farmer)
 
     # manually add pending request
-    farmer.pending_solver_requests[partial_proofs.partial_proofs[0]] = {
+    key = bytes(partial_proofs.partial_proofs[0])
+    farmer.pending_solver_requests[key] = {
         "proof_data": partial_proofs,
         "peer": harvester_peer,
     }
@@ -477,7 +486,8 @@ async def test_solution_response_handler(
         assert original_peer == harvester_peer
 
         # verify pending request was removed
-        assert partial_proofs.partial_proofs[0] not in farmer.pending_solver_requests
+        key = bytes(partial_proofs.partial_proofs[0])
+        assert key not in farmer.pending_solver_requests
 
 
 @pytest.mark.anyio
@@ -492,7 +502,9 @@ async def test_solution_response_unknown_quality(
     solver_peer = await get_solver_peer(farmer)
 
     # create solution response with unknown quality
-    solution_response = solver_protocol.SolverResponse(partial_proof=bytes(b"1" * 32), proof=b"test_proof")
+    solution_response = solver_protocol.SolverResponse(
+        partial_proof=[uint64(1), uint64(2), uint64(3), uint64(4)], proof=b"test_proof"
+    )
 
     with unittest.mock.patch.object(farmer_api, "new_proof_of_space", new_callable=AsyncMock) as mock_new_proof:
         await farmer_api.solution_response(solution_response, solver_peer)
@@ -518,9 +530,11 @@ async def test_solution_response_empty_proof(
         challenge_hash=challenge_hash,
         sp_hash=sp_hash,
         plot_identifier="test_plot_id",
-        partial_proofs=[b"test_partial_proof_for_quality"],
+        partial_proofs=[[uint64(1), uint64(2), uint64(3), uint64(4)], [uint64(2), uint64(3), uint64(4), uint64(5)]],
         signage_point_index=uint8(0),
         plot_size=uint8(32),
+        plot_id=bytes32.fromhex("abababababababababababababababababababababababababababababababab"),
+        strength=uint8(5),
         pool_public_key=G1Element(),
         pool_contract_puzzle_hash=bytes32(b"4" * 32),
         plot_public_key=G1Element(),
@@ -530,8 +544,9 @@ async def test_solution_response_empty_proof(
     harvester_peer.peer_node_id = "harvester_peer"
 
     # manually add pending request
-    farmer.pending_solver_requests[partial_proofs.partial_proofs[0]] = {
-        "proof_data": partial_proofs.partial_proofs[0],
+    key = bytes(partial_proofs.partial_proofs[0])
+    farmer.pending_solver_requests[key] = {
+        "proof_data": partial_proofs,
         "peer": harvester_peer,
     }
 
@@ -548,7 +563,8 @@ async def test_solution_response_empty_proof(
         mock_new_proof.assert_not_called()
 
         # verify pending request was removed (cleanup still happens)
-        assert partial_proofs.partial_proofs[0] not in farmer.pending_solver_requests
+        key = bytes(partial_proofs.partial_proofs[0])
+        assert key not in farmer.pending_solver_requests
 
 
 @pytest.mark.anyio
@@ -579,9 +595,11 @@ async def test_v2_partial_proofs_solver_exception(
         challenge_hash=challenge_hash,
         sp_hash=sp_hash,
         plot_identifier="test_plot_id",
-        partial_proofs=[b"test_partial_proof_1"],
+        partial_proofs=[[uint64(1), uint64(2), uint64(3), uint64(4)], [uint64(2), uint64(3), uint64(4), uint64(5)]],
         signage_point_index=uint8(0),
         plot_size=uint8(32),
+        plot_id=bytes32.fromhex("abababababababababababababababababababababababababababababababab"),
+        strength=uint8(5),
         pool_public_key=G1Element(),
         pool_contract_puzzle_hash=bytes32(b"4" * 32),
         plot_public_key=G1Element(),
@@ -594,4 +612,5 @@ async def test_v2_partial_proofs_solver_exception(
         await farmer_api.partial_proofs(partial_proofs, harvester_peer)
 
         # verify pending request was cleaned up after exception
-        assert partial_proofs.partial_proofs[0] not in farmer.pending_solver_requests
+        key = bytes(partial_proofs.partial_proofs[0])
+        assert key not in farmer.pending_solver_requests

chia/_tests/harvester/test_harvester_api.py

@@ -15,6 +15,7 @@
 from chia._tests.plotting.util import get_test_plots
 from chia._tests.util.time_out_assert import time_out_assert
 from chia.harvester.harvester_api import HarvesterAPI
+from chia.plotting.prover import V1Prover, V2Prover
 from chia.plotting.util import PlotInfo
 from chia.protocols import harvester_protocol
 from chia.protocols.harvester_protocol import PoolDifficulty
@@ -85,10 +86,16 @@ def create_test_setup(
 
 @contextmanager
 def mock_successful_proof(plot_info: PlotInfo) -> Iterator[None]:
-    with patch.object(plot_info.prover, "get_full_proof") as mock_get_proof:
-        mock_proof = MagicMock(spec=ProofOfSpace)
-        mock_get_proof.return_value = mock_proof, None
-        yield
+    if isinstance(plot_info.prover, V1Prover):
+        with patch.object(plot_info.prover, "get_full_proof") as mock_get_proof:
+            mock_proof = MagicMock(spec=ProofOfSpace)
+            mock_get_proof.return_value = mock_proof, None
+            yield
+    elif isinstance(plot_info.prover, V2Prover):
+        with patch.object(plot_info.prover, "get_partial_proof") as mock_get_proof:
+            mock_proof = MagicMock(spec=ProofOfSpace)
+            mock_get_proof.return_value = [uint64(1)] * 64, None
+            yield
 
 
 def assert_farming_info_sent(mock_peer: MagicMock) -> None:

chia/_tests/plotting/test_prover.py

@@ -12,44 +12,57 @@
 
 
 class TestProver:
+    def test_get_prover_from_file_with_plot1_still_works(self) -> None:
+        with tempfile.NamedTemporaryFile(suffix=".plot", delete=False) as f:
+            temp_path = f.name
+        try:
+            with pytest.raises(Exception) as exc_info:
+                get_prover_from_file(temp_path)
+            assert not isinstance(exc_info.value, NotImplementedError)
+        finally:
+            Path(temp_path).unlink()
+
+    def test_unsupported_file_extension_raises_value_error(self) -> None:
+        with pytest.raises(ValueError, match="Unsupported plot file"):
+            get_prover_from_file("/nonexistent/path/test.txt")
+
+
+# TODO: todo_v2_plots enable these tests once we have test plots we can load
+@pytest.mark.skip("we don't have v2 test plots yet")
+class TestV2Prover:
     def test_v2_prover_init_with_nonexistent_file(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
+        prover = V2Prover.from_filename("/nonexistent/path/test.plot2")
         assert prover.get_version() == PlotVersion.V2
         assert prover.get_filename() == "/nonexistent/path/test.plot2"
 
     def test_v2_prover_get_size_raises_error(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
+        prover = V2Prover.from_filename("/nonexistent/path/test.plot2")
         with pytest.raises(NotImplementedError, match="V2 plot format is not yet implemented"):
             prover.get_size()
 
     def test_v2_prover_get_memo_raises_error(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
+        prover = V2Prover.from_filename("/nonexistent/path/test.plot2")
         with pytest.raises(NotImplementedError, match="V2 plot format is not yet implemented"):
             prover.get_memo()
 
     def test_v2_prover_get_compression_level(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
+        prover = V2Prover.from_filename("/nonexistent/path/test.plot2")
         assert prover.get_compression_level() == uint8(0)
 
     def test_v2_prover_get_id_raises_error(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
+        prover = V2Prover.from_filename("/nonexistent/path/test.plot2")
         with pytest.raises(NotImplementedError, match="V2 plot format is not yet implemented"):
             prover.get_id()
 
     def test_v2_prover_get_qualities_for_challenge_raises_error(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
+        prover = V2Prover.from_filename("/nonexistent/path/test.plot2")
         with pytest.raises(
             AssertionError, match="V2 plot format does not support qualities directly, use partial proofs"
         ):
-            prover.get_qualities_for_challenge(bytes32(b"1" * 32))
-
-    def test_v2_prover_get_full_proof_raises_error(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
-        with pytest.raises(AssertionError, match="V2 plot format require solver to get full proof"):
-            prover.get_full_proof(bytes32(b"1" * 32), 0)
+            prover.get_qualities_for_challenge(bytes32(b"1" * 32), uint8(5))
 
     def test_v2_prover_bytes_raises_error(self) -> None:
-        prover = V2Prover("/nonexistent/path/test.plot2")
+        prover = V2Prover.from_filename("/nonexistent/path/test.plot2")
         with pytest.raises(NotImplementedError, match="V2 plot format is not yet implemented"):
             bytes(prover)
 
@@ -63,20 +76,6 @@ def test_get_prover_from_file(self) -> None:
         with pytest.raises(NotImplementedError, match="V2 plot format is not yet implemented"):
             prover.get_size()
 
-    def test_get_prover_from_file_with_plot1_still_works(self) -> None:
-        with tempfile.NamedTemporaryFile(suffix=".plot", delete=False) as f:
-            temp_path = f.name
-        try:
-            with pytest.raises(Exception) as exc_info:
-                get_prover_from_file(temp_path)
-            assert not isinstance(exc_info.value, NotImplementedError)
-        finally:
-            Path(temp_path).unlink()
-
-    def test_unsupported_file_extension_raises_value_error(self) -> None:
-        with pytest.raises(ValueError, match="Unsupported plot file"):
-            get_prover_from_file("/nonexistent/path/test.txt")
-
 
 class TestV1Prover:
     def test_v1_prover_get_version(self) -> None:
@@ -87,13 +86,6 @@ def test_v1_prover_get_version(self) -> None:
 
 
 class TestGetProverFromBytes:
-    def test_get_prover_from_bytes_v2_plot(self) -> None:
-        with patch("chia.plotting.prover.V2Prover.from_bytes") as mock_v2_from_bytes:
-            mock_prover = MagicMock()
-            mock_v2_from_bytes.return_value = mock_prover
-            result = get_prover_from_bytes("test.plot2", b"test_data")
-            assert result == mock_prover
-
     def test_get_prover_from_bytes_v1_plot(self) -> None:
         with patch("chia.plotting.prover.DiskProver") as mock_disk_prover_class:
             mock_disk_prover = MagicMock()
@@ -102,5 +94,5 @@ def test_get_prover_from_bytes_v1_plot(self) -> None:
             assert isinstance(result, V1Prover)
 
     def test_get_prover_from_bytes_unsupported_extension(self) -> None:
-        with pytest.raises(ValueError, match="Unsupported plot file"):
+        with pytest.raises((RuntimeError, ValueError)):
             get_prover_from_bytes("test.txt", b"test_data")

chia/_tests/solver/test_solver_service.py

@@ -5,6 +5,8 @@
 
 import pytest
 from chia_rs import ConsensusConstants
+from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint8, uint64
 
 from chia.protocols.outbound_message import Message
 from chia.protocols.solver_protocol import SolverInfo
@@ -21,7 +23,12 @@ async def test_solver_api_methods(blockchain_constants: ConsensusConstants, tmp_
             solver = solver_service._node
             solver_api = solver_service._api
             assert solver_api.ready() is True
-            test_info = SolverInfo(partial_proof=b"test_partial_proof_42")
+            test_info = SolverInfo(
+                partial_proof=[uint64(1), uint64(2), uint64(3), uint64(4)],
+                plot_id=bytes32.fromhex("abababababababababababababababababababababababababababababababab"),
+                strength=uint8(5),
+                size=uint8(28),
+            )
             expected_proof = b"test_proof_data_12345"
             with patch.object(solver, "solve", return_value=expected_proof):
                 api_result = await solver_api.solve(test_info)