From 5cf019ff5ecf7757e6670cf9a9a3698f519dac4c Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 20 Jan 2025 11:37:23 +0200
Subject: [PATCH] pass augmented in tests

---
 chia/_tests/util/full_sync.py               | 3 +++
 chia/_tests/wallet/sync/test_wallet_sync.py | 5 +++++
 tools/test_full_sync.py                     | 4 ++++
 3 files changed, 12 insertions(+)

diff --git a/chia/_tests/util/full_sync.py b/chia/_tests/util/full_sync.py
index 362afbbbfb0f..4d07f41d386b 100644
--- a/chia/_tests/util/full_sync.py
+++ b/chia/_tests/util/full_sync.py
@@ -28,6 +28,7 @@
 from chia.types.full_block import FullBlock
 from chia.types.peer_info import PeerInfo
 from chia.types.validation_state import ValidationState
+from chia.util.augmented_chain import AugmentedBlockchain
 from chia.util.config import load_config
 from chia.util.ints import uint16
 
@@ -184,6 +185,7 @@ async def run_sync_test(
                 worst_batch_height = None
                 worst_batch_time_per_block = None
                 peer_info = peer.get_peer_logging()
+                blockchain = AugmentedBlockchain(full_node.blockchain)
                 async for r in rows:
                     batch_start_time = time.monotonic()
                     with enable_profiler(profile, height):
@@ -216,6 +218,7 @@ async def run_sync_test(
                                 peer_info,
                                 ForkInfo(fork_height, fork_height, header_hash),
                                 ValidationState(ssi, diff, None),
+                                blockchain
                             )
                             end_height = block_batch[-1].height
                             full_node.blockchain.clean_block_record(end_height - full_node.constants.BLOCKS_CACHE_SIZE)
diff --git a/chia/_tests/wallet/sync/test_wallet_sync.py b/chia/_tests/wallet/sync/test_wallet_sync.py
index c9737b165fc6..92dea1984dda 100644
--- a/chia/_tests/wallet/sync/test_wallet_sync.py
+++ b/chia/_tests/wallet/sync/test_wallet_sync.py
@@ -55,6 +55,7 @@
 from chia.types.peer_info import PeerInfo
 from chia.types.spend_bundle import SpendBundle
 from chia.types.validation_state import ValidationState
+from chia.util.augmented_chain import AugmentedBlockchain
 from chia.util.hash import std_hash
 from chia.util.ints import uint32, uint64, uint128
 from chia.wallet.nft_wallet.nft_wallet import NFTWallet
@@ -431,11 +432,13 @@ async def test_long_sync_wallet(
         full_node.constants, True, block_record, full_node.blockchain
     )
     fork_height = blocks_reorg[-num_blocks - 10].height - 1
+    blockchain = AugmentedBlockchain(full_node.blockchain)
     await full_node.add_block_batch(
         blocks_reorg[-num_blocks - 10 : -1],
         PeerInfo("0.0.0.0", 0),
         ForkInfo(fork_height, fork_height, blocks_reorg[-num_blocks - 10].prev_header_hash),
         ValidationState(sub_slot_iters, difficulty, None),
+        blockchain,
     )
     await full_node.add_block(blocks_reorg[-1])
 
@@ -559,11 +562,13 @@ async def test_wallet_reorg_get_coinbase(
     sub_slot_iters, difficulty = get_next_sub_slot_iters_and_difficulty(
         full_node.constants, True, block_record, full_node.blockchain
     )
+    blockchain = AugmentedBlockchain(full_node.blockchain)
     await full_node.add_block_batch(
         blocks_reorg_2[-44:],
         PeerInfo("0.0.0.0", 0),
         ForkInfo(blocks_reorg_2[-45].height, blocks_reorg_2[-45].height, blocks_reorg_2[-45].header_hash),
         ValidationState(sub_slot_iters, difficulty, None),
+        blockchain,
     )
 
     for wallet_node, wallet_server in wallets:
diff --git a/tools/test_full_sync.py b/tools/test_full_sync.py
index 0abb7db4efd8..b19e7cbfcc3c 100755
--- a/tools/test_full_sync.py
+++ b/tools/test_full_sync.py
@@ -21,6 +21,7 @@
 from chia.server.ws_connection import WSChiaConnection
 from chia.types.full_block import FullBlock
 from chia.types.validation_state import ValidationState
+from chia.util.augmented_chain import AugmentedBlockchain
 from chia.util.config import load_config
 
 
@@ -149,6 +150,7 @@ async def run_sync_checkpoint(
 
             block_batch = []
             peer_info = peer.get_peer_logging()
+            blockchain = AugmentedBlockchain(full_node.blockchain)
             async for r in rows:
                 block = FullBlock.from_bytes_unchecked(zstd.decompress(r[0]))
                 block_batch.append(block)
@@ -169,6 +171,7 @@ async def run_sync_checkpoint(
                     peer_info,
                     ForkInfo(fork_height, fork_height, header_hash),
                     ValidationState(ssi, diff, None),
+                    blockchain
                 )
                 end_height = block_batch[-1].height
                 full_node.blockchain.clean_block_record(end_height - full_node.constants.BLOCKS_CACHE_SIZE)
@@ -192,6 +195,7 @@ async def run_sync_checkpoint(
                     peer_info,
                     ForkInfo(fork_height, fork_height, fork_header_hash),
                     ValidationState(ssi, diff, None),
+                    blockchain
                 )
                 if not success:
                     raise RuntimeError("failed to ingest block batch")