more

Author: richardkiss

chia/_tests/blockchain/test_blockchain.py

@@ -3886,7 +3886,8 @@ async def test_chain_failed_rollback(empty_blockchain: Blockchain, bt: BlockTool
     print(f"{recs_dbg1[0] if len(recs_dbg1) > 0 else None}")
     print(spend_bundle.coin_spends[0].coin.name())
     # await b.consensus_store._set_spent([spend_bundle.coin_spends[0].coin.name()], 8)
-    await b.consensus_store.rollback_to_block(2)
+    async with b.consensus_store as writer:
+        await writer.rollback_to_block(2)
     recs_dbg2 = await b.consensus_store.get_coin_records([spend_bundle.coin_spends[0].coin.name()])
     print(f"{recs_dbg2[0] if len(recs_dbg2) > 0 else None}")
 

chia/consensus/consensus_store_protocol.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Collection
 from typing import AsyncIterator, Optional, Protocol
+from types import TracebackType
 
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
@@ -61,7 +62,12 @@ class ConsensusStoreProtocol(Protocol):
 
     # Async context manager methods
     async def __aenter__(self) -> ConsensusStoreWriteProtocol: ...
-    async def __aexit__(self, exc_type, exc, tb) -> Optional[bool]: ...
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> Optional[bool]: ...
 
     # Block store reads
     async def get_block_records_close_to_peak(

chia/full_node/consensus_store_sqlite3.py

@@ -4,6 +4,7 @@
 from collections.abc import Collection
 from contextlib import AbstractAsyncContextManager
 from typing import Any, AsyncIterator, Optional, TYPE_CHECKING
+from types import TracebackType
 
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
@@ -97,7 +98,12 @@ async def __aenter__(self):
         self._txn_depth += 1
         return self._writer  # Return the same writer for nested contexts
 
-    async def __aexit__(self, exc_type, exc, tb):
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> Optional[bool]:
         try:
             # Check if we're at the outermost level before decrementing
             if self._txn_depth == 1: