From 533bacb2348eaaaf964719b00655f9780c80d781 Mon Sep 17 00:00:00 2001
From: MiniFrenchBread <103425574+MiniFrenchBread@users.noreply.github.com>
Date: Tue, 23 Jul 2024 15:47:44 +0800
Subject: [PATCH] fix: admin_getFileLocation; test: sync test (#141)

* feat: add all_shards in admin_getFileLocation

* fix: admin_getFileLocation

* test: improve sync test

* fix: lint
---
 node/rpc/src/admin/impl.rs       |  6 +++++-
 tests/sync_test.py               | 13 +++++++++++--
 tests/test_framework/zgs_node.py |  3 +++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/node/rpc/src/admin/impl.rs b/node/rpc/src/admin/impl.rs
index 8ec5b4f..db6b0ff 100644
--- a/node/rpc/src/admin/impl.rs
+++ b/node/rpc/src/admin/impl.rs
@@ -190,6 +190,8 @@ impl RpcServer for RpcServerImpl {
         tx_seq: u64,
         all_shards: bool,
     ) -> RpcResult<Option<Vec<LocationInfo>>> {
+        info!("admin_getFileLocation()");
+
         let tx = match self.ctx.log_store.get_tx_by_seq_number(tx_seq).await? {
             Some(tx) => tx,
             None => {
@@ -225,7 +227,9 @@ impl RpcServer for RpcServerImpl {
                 shard_config: shard_config.unwrap(),
             })
             .collect();
-        if all_shards && all_shards_available(info.iter().map(|info| info.shard_config).collect()) {
+
+        if !all_shards || all_shards_available(info.iter().map(|info| info.shard_config).collect())
+        {
             Ok(Some(info))
         } else {
             Ok(None)
diff --git a/tests/sync_test.py b/tests/sync_test.py
index cb16fdf..074b01f 100755
--- a/tests/sync_test.py
+++ b/tests/sync_test.py
@@ -29,6 +29,9 @@ class SyncTest(TestFramework):
         client1 = self.nodes[0]
         client2 = self.nodes[1]
 
+        # stop client2, preventing it from receiving AnnounceFile
+        client2.shutdown()
+
         # Create submission
         chunk_data = random.randbytes(256 * 1024)
         data_root = self.__create_submission(chunk_data)
@@ -41,16 +44,22 @@ class SyncTest(TestFramework):
         segments = submit_data(client1, chunk_data)
         self.log.info("segments: %s", [(s["root"], s["index"], s["proof"]) for s in segments])
         wait_until(lambda: client1.zgs_get_file_info(data_root)["finalized"])
-
-        # File should not be auto sync on node 2
+        
+        # restart client2
+        client2.start()
+        client2.wait_for_rpc_connection()
+        
+        # File should not be auto sync on node 2 and there is no cached file locations
         wait_until(lambda: client2.zgs_get_file_info(data_root) is not None)
         time.sleep(3)
         assert_equal(client2.zgs_get_file_info(data_root)["finalized"], False)
+        assert(client2.admin_get_file_location(0) is None)
 
         # Trigger file sync by rpc
         assert(client2.admin_start_sync_file(0) is None)
         wait_until(lambda: client2.sync_status_is_completed_or_unknown(0))
         wait_until(lambda: client2.zgs_get_file_info(data_root)["finalized"])
+        assert(client2.admin_get_file_location(0) is not None)
 
         # Validate data
         assert_equal(
diff --git a/tests/test_framework/zgs_node.py b/tests/test_framework/zgs_node.py
index 0fe7caf..19a4256 100644
--- a/tests/test_framework/zgs_node.py
+++ b/tests/test_framework/zgs_node.py
@@ -113,6 +113,9 @@ class ZgsNode(TestNode):
     def sync_status_is_completed_or_unknown(self, tx_seq):
         status = self.rpc.admin_getSyncStatus([tx_seq])
         return status == "Completed" or status == "unknown"
+    
+    def admin_get_file_location(self, tx_seq, all_shards = True):
+        return self.rpc.admin_getFileLocation([tx_seq, all_shards])
 
     def clean_data(self):
         shutil.rmtree(os.path.join(self.data_dir, "db"))