From b2a70501c2a36161e4a86e6dce84bd893ed76952 Mon Sep 17 00:00:00 2001
From: bruno-valante <140794260+bruno-valante@users.noreply.github.com>
Date: Tue, 15 Oct 2024 14:24:56 +0800
Subject: [PATCH] Test flow root consistency (#230)

---
 node/rpc/src/zgs/api.rs                |  5 ++-
 node/rpc/src/zgs/impl.rs               |  6 ++-
 tests/root_consistency_test.py         | 55 ++++++++++++++++++++++++++
 tests/test_framework/contract_proxy.py |  5 +++
 tests/test_framework/zgs_node.py       |  3 ++
 5 files changed, 72 insertions(+), 2 deletions(-)
 create mode 100755 tests/root_consistency_test.py

diff --git a/node/rpc/src/zgs/api.rs b/node/rpc/src/zgs/api.rs
index 9cc07e8..32dcfb4 100644
--- a/node/rpc/src/zgs/api.rs
+++ b/node/rpc/src/zgs/api.rs
@@ -2,7 +2,7 @@ use crate::types::{FileInfo, Segment, SegmentWithProof, Status};
 use jsonrpsee::core::RpcResult;
 use jsonrpsee::proc_macros::rpc;
 use shared_types::{DataRoot, FlowProof, TxSeqOrRoot};
-use storage::config::ShardConfig;
+use storage::{config::ShardConfig, H256};
 
 #[rpc(server, client, namespace = "zgs")]
 pub trait Rpc {
@@ -77,4 +77,7 @@ pub trait Rpc {
         sector_index: u64,
         flow_root: Option<DataRoot>,
     ) -> RpcResult<FlowProof>;
+
+    #[method(name = "getFlowContext")]
+    async fn get_flow_context(&self) -> RpcResult<(H256, u64)>;
 }
diff --git a/node/rpc/src/zgs/impl.rs b/node/rpc/src/zgs/impl.rs
index 4c1e0f1..e9f6295 100644
--- a/node/rpc/src/zgs/impl.rs
+++ b/node/rpc/src/zgs/impl.rs
@@ -8,7 +8,7 @@ use jsonrpsee::core::RpcResult;
 use shared_types::{DataRoot, FlowProof, Transaction, TxSeqOrRoot, CHUNK_SIZE};
 use std::fmt::{Debug, Formatter, Result};
 use storage::config::ShardConfig;
-use storage::try_option;
+use storage::{try_option, H256};
 
 pub struct RpcServerImpl {
     pub ctx: Context,
@@ -198,6 +198,10 @@ impl RpcServer for RpcServerImpl {
         assert_eq!(proof.left_proof, proof.right_proof);
         Ok(proof.right_proof)
     }
+
+    async fn get_flow_context(&self) -> RpcResult<(H256, u64)> {
+        Ok(self.ctx.log_store.get_context().await?)
+    }
 }
 
 impl RpcServerImpl {
diff --git a/tests/root_consistency_test.py b/tests/root_consistency_test.py
new file mode 100755
index 0000000..737088b
--- /dev/null
+++ b/tests/root_consistency_test.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+from test_framework.test_framework import TestFramework
+from config.node_config import MINER_ID, GENESIS_PRIV_KEY
+from utility.submission import create_submission, submit_data
+from utility.utils import wait_until, assert_equal
+from test_framework.blockchain_node import BlockChainNodeType
+
+
+class RootConsistencyTest(TestFramework):
+    def setup_params(self):
+        self.num_blockchain_nodes = 1
+        self.num_nodes = 1
+
+    def submit_data(self, item, size):
+        submissions_before = self.contract.num_submissions()
+        client = self.nodes[0]
+        chunk_data = item * 256 * size
+        submissions, data_root = create_submission(chunk_data)
+        self.contract.submit(submissions)
+        wait_until(lambda: self.contract.num_submissions() == submissions_before + 1)
+        wait_until(lambda: client.zgs_get_file_info(data_root) is not None)
+
+        segment = submit_data(client, chunk_data)
+        wait_until(lambda: client.zgs_get_file_info(data_root)["finalized"])
+
+    def assert_flow_status(self, expected_length):
+        contract_root = self.contract.get_flow_root().hex()
+        contract_length = self.contract.get_flow_length()
+        (node_root, node_length) = tuple(self.nodes[0].zgs_getFlowContext())
+
+        assert_equal(contract_length, node_length)
+        assert_equal(contract_length, expected_length)
+        assert_equal(contract_root, node_root[2:])
+
+
+
+    def run_test(self):
+        self.assert_flow_status(1)
+
+        self.submit_data(b"\x11", 1)
+        self.assert_flow_status(2)
+
+        self.submit_data(b"\x11", 8 + 4 + 2)
+        self.assert_flow_status(16 + 4 + 2)
+
+        self.submit_data(b"\x12", 128 + 64)
+        self.assert_flow_status(256 + 64)
+
+        self.submit_data(b"\x13", 512 + 256)
+        self.assert_flow_status(1024 + 512 + 256)
+        
+
+
+if __name__ == "__main__":
+    RootConsistencyTest().main()
diff --git a/tests/test_framework/contract_proxy.py b/tests/test_framework/contract_proxy.py
index a9e05f1..83c9ebe 100644
--- a/tests/test_framework/contract_proxy.py
+++ b/tests/test_framework/contract_proxy.py
@@ -91,7 +91,12 @@ class FlowContractProxy(ContractProxy):
 
     def get_mine_context(self, node_idx=0):
         return self._call("makeContextWithResult", node_idx)
+    
+    def get_flow_root(self, node_idx=0):
+        return self._call("computeFlowRoot", node_idx)
 
+    def get_flow_length(self, node_idx=0):
+        return self._call("tree", node_idx)[0]
 
 class MineContractProxy(ContractProxy):
     def last_mined_epoch(self, node_idx=0):
diff --git a/tests/test_framework/zgs_node.py b/tests/test_framework/zgs_node.py
index 146799f..a27a194 100644
--- a/tests/test_framework/zgs_node.py
+++ b/tests/test_framework/zgs_node.py
@@ -100,6 +100,9 @@ class ZgsNode(TestNode):
 
     def zgs_get_file_info_by_tx_seq(self, tx_seq):
         return self.rpc.zgs_getFileInfoByTxSeq([tx_seq])
+    
+    def zgs_get_flow_context(self, tx_seq):
+        return self.rpc.zgs_getFlowContext([tx_seq])
 
     def shutdown(self):
         self.rpc.admin_shutdown()