Undo test changes and silly api

Author: richardkiss

chia/_tests/blockchain/blockchain_test_utils.py

@@ -10,25 +10,42 @@
 from chia.consensus.blockchain import AddBlockResult, Blockchain
 from chia.consensus.difficulty_adjustment import get_next_sub_slot_iters_and_difficulty
 from chia.consensus.multiprocess_validation import PreValidationResult, pre_validate_block
+from chia.full_node.consensus_store_sqlite3 import ConsensusStoreSQLite3
 from chia.types.validation_state import ValidationState
 from chia.util.errors import Err
 
 
 async def check_block_store_invariant(bc: Blockchain):
+    # This function checks the invariant of the sqlite database.
+    # Only operate on the sqlite block store.
+    if not isinstance(bc.consensus_store, ConsensusStoreSQLite3):
+        return
+
+    db_wrapper = bc.consensus_store.block_store.db_wrapper
+
+    if db_wrapper.db_version == 1:
+        return
+
     in_chain = set()
     max_height = -1
-    async for height in bc.consensus_store.get_block_heights_in_main_chain():
-        # if this block is in-chain, ensure we haven't found another block
-        # at this height that's also in chain. That would be an invariant
-        # violation
-        # make sure we don't have any duplicate heights. Each block
-        # height can only have a single block with in_main_chain set
-        assert height not in in_chain
-        in_chain.add(height)
-        max_height = max(max_height, height)
-
-    # make sure every height is represented in the set
-    assert len(in_chain) == max_height + 1
+    async with bc.consensus_store.block_store.transaction() as conn:
+        async with conn.execute("SELECT height, in_main_chain FROM full_blocks") as cursor:
+            rows = await cursor.fetchall()
+            for row in rows:
+                height = row[0]
+
+                # if this block is in-chain, ensure we haven't found another block
+                # at this height that's also in chain. That would be an invariant
+                # violation
+                if row[1]:
+                    # make sure we don't have any duplicate heights. Each block
+                    # height can only have a single block with in_main_chain set
+                    assert height not in in_chain
+                    in_chain.add(height)
+                    max_height = max(max_height, height)
+
+            # make sure every height is represented in the set
+            assert len(in_chain) == max_height + 1
 
 
 async def _validate_and_add_block(

chia/consensus/consensus_store_protocol.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Collection
+from collections.abc import Collection
 from types import TracebackType
 from typing import Optional, Protocol
 
@@ -85,7 +85,6 @@ async def get_sub_epoch_challenge_segments(
     ) -> Optional[list[SubEpochChallengeSegment]]: ...
     async def get_generator(self, header_hash: bytes32) -> Optional[bytes]: ...
     async def get_generators_at(self, heights: set[uint32]) -> dict[uint32, bytes]: ...
-    def get_block_heights_in_main_chain(self) -> AsyncIterator[int]: ...
 
     # Coin store reads
     async def get_coin_records(self, names: Collection[bytes32]) -> list[CoinRecord]: ...

chia/full_node/consensus_store_sqlite3.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Collection
+from collections.abc import Collection
 from contextlib import AbstractAsyncContextManager
 from dataclasses import dataclass
 from pathlib import Path
@@ -187,16 +187,6 @@ async def get_coins_added_at_height(self, height: uint32) -> list[CoinRecord]:
     async def get_coins_removed_at_height(self, height: uint32) -> list[CoinRecord]:
         return await self.coin_store.get_coins_removed_at_height(height)
 
-    def get_block_heights_in_main_chain(self) -> AsyncIterator[int]:
-        async def gen() -> AsyncIterator[int]:
-            async with self.block_store.transaction() as conn:
-                async with conn.execute("SELECT height, in_main_chain FROM full_blocks") as cursor:
-                    async for row in cursor:
-                        if row[1]:
-                            yield row[0]
-
-        return gen()
-
     # Height map methods
     def get_ses_heights(self) -> list[uint32]:
         return self.height_map.get_ses_heights()