From fd1b7c1967915ec56a2cc7d86833f6f5e04ca71c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adam=20Chuda=C5=9B?= <adam@chudas.net>
Date: Mon, 9 Dec 2024 13:18:47 +0100
Subject: [PATCH] update tests

---
 .../src/tests/client/cold_storage.rs          | 27 +++++++++++++------
 .../src/tests/client/process_blocks.rs        | 12 ++++++++-
 tools/cold-store/src/cli.rs                   | 18 +++++++------
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/integration-tests/src/tests/client/cold_storage.rs b/integration-tests/src/tests/client/cold_storage.rs
index 64e05437bf2..0512eb602b4 100644
--- a/integration-tests/src/tests/client/cold_storage.rs
+++ b/integration-tests/src/tests/client/cold_storage.rs
@@ -147,9 +147,12 @@ fn test_storage_after_commit_of_cold_update() {
 
         let client = &env.clients[0];
         let client_store = client.runtime_adapter.store();
-        let epoch_id = client.epoch_manager.get_epoch_id_from_prev_block(&last_hash).unwrap();
+        let epoch_id = client.epoch_manager.get_epoch_id(block.hash()).unwrap();
         let shard_layout = client.epoch_manager.get_shard_layout(&epoch_id).unwrap();
-        update_cold_db(cold_db, &client_store, &shard_layout, &height, 4).unwrap();
+        let is_last_block_in_epoch =
+            client.epoch_manager.is_next_block_epoch_start(block.hash()).unwrap();
+        update_cold_db(cold_db, &client_store, &shard_layout, &height, is_last_block_in_epoch, 4)
+            .unwrap();
 
         last_hash = *block.hash();
     }
@@ -281,14 +284,22 @@ fn test_cold_db_copy_with_height_skips() {
         };
 
         let client = &env.clients[0];
-        let epoch_id = client.epoch_manager.get_epoch_id_from_prev_block(&last_hash).unwrap();
+        let hot_store = client.runtime_adapter.store();
+        let block_hash =
+            hot_store.get_ser::<CryptoHash>(DBCol::BlockHeight, &height.to_le_bytes()).unwrap();
+        let Some(block) = block else {
+            assert!(block_hash.is_none());
+            continue;
+        };
+        let block_hash = block_hash.unwrap();
+        assert!(&block_hash == block.hash());
+        let epoch_id = client.epoch_manager.get_epoch_id(&block_hash).unwrap();
         let shard_layout = client.epoch_manager.get_shard_layout(&epoch_id).unwrap();
-        update_cold_db(&cold_db, &client.runtime_adapter.store(), &shard_layout, &height, 1)
+        let is_last_block_in_epoch =
+            client.epoch_manager.is_next_block_epoch_start(&block_hash).unwrap();
+        update_cold_db(&cold_db, hot_store, &shard_layout, &height, is_last_block_in_epoch, 1)
             .unwrap();
-
-        if block.is_some() {
-            last_hash = *block.unwrap().hash();
-        }
+        last_hash = block_hash;
     }
 
     // We still need to filter out one chunk
diff --git a/integration-tests/src/tests/client/process_blocks.rs b/integration-tests/src/tests/client/process_blocks.rs
index 2f97cd0f100..7ced7a4c77b 100644
--- a/integration-tests/src/tests/client/process_blocks.rs
+++ b/integration-tests/src/tests/client/process_blocks.rs
@@ -1464,11 +1464,21 @@ fn test_archival_gc_common(
         let header = block.header();
         let epoch_id = header.epoch_id();
         let shard_layout = env.clients[0].epoch_manager.get_shard_layout(epoch_id).unwrap();
+        let is_last_block_in_epoch =
+            env.clients[0].epoch_manager.is_next_block_epoch_start(header.hash()).unwrap();
 
         blocks.push(block);
 
         if i <= max_cold_head_height {
-            update_cold_db(storage.cold_db().unwrap(), hot_store, &shard_layout, &i, 1).unwrap();
+            update_cold_db(
+                storage.cold_db().unwrap(),
+                hot_store,
+                &shard_layout,
+                &i,
+                is_last_block_in_epoch,
+                1,
+            )
+            .unwrap();
             update_cold_head(storage.cold_db().unwrap(), &hot_store, &i).unwrap();
         }
     }
diff --git a/tools/cold-store/src/cli.rs b/tools/cold-store/src/cli.rs
index 651ab4069ab..7fc2c99ea50 100644
--- a/tools/cold-store/src/cli.rs
+++ b/tools/cold-store/src/cli.rs
@@ -215,24 +215,26 @@ fn copy_next_block(store: &NodeStorage, config: &NearConfig, epoch_manager: &Epo
 
     // Here it should be sufficient to just read from hot storage.
     // Because BlockHeight is never garbage collectable and is not even copied to cold.
-    let cold_head_hash = get_ser_from_store::<CryptoHash>(
+    let next_height_block_hash = get_ser_from_store::<CryptoHash>(
         &store.get_hot_store(),
         DBCol::BlockHeight,
-        &cold_head_height.to_le_bytes(),
+        &next_height.to_le_bytes(),
     )
-    .unwrap_or_else(|| panic!("No block hash in hot storage for height {}", cold_head_height));
+    .unwrap_or_else(|| panic!("No block hash in hot storage for height {}", next_height));
 
     // For copying block we need to have shard_layout.
     // For that we need epoch_id.
-    // For that we might use prev_block_hash, and because next_hight = cold_head_height + 1,
-    // we use cold_head_hash.
+    // For that we might use the hash of the block.
+    let epoch_id = &epoch_manager.get_epoch_id(&next_height_block_hash).unwrap();
+    let shard_layout = &epoch_manager.get_shard_layout(epoch_id).unwrap();
+    let is_last_block_in_epoch =
+        epoch_manager.is_next_block_epoch_start(&next_height_block_hash).unwrap();
     update_cold_db(
         &*store.cold_db().unwrap(),
         &store.get_hot_store(),
-        &epoch_manager
-            .get_shard_layout(&epoch_manager.get_epoch_id_from_prev_block(&cold_head_hash).unwrap())
-            .unwrap(),
+        shard_layout,
         &next_height,
+        is_last_block_in_epoch,
         1,
     )
     .unwrap_or_else(|_| panic!("Failed to copy block at height {} to cold db", next_height));