From 9952d4eb364e54f0d0e11c5ffb057c49f62027a7 Mon Sep 17 00:00:00 2001
From: Aliaksandr Tsurko <ales@opentensor.dev>
Date: Wed, 30 Oct 2024 19:53:23 +0100
Subject: [PATCH] Add test for stake value consistency

---
 pallets/subtensor/src/rpc_info/neuron_info.rs | 12 +++++++++
 pallets/subtensor/src/rpc_info/stake_info.rs  |  6 +++++
 pallets/subtensor/tests/staking.rs            | 25 +++++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/pallets/subtensor/src/rpc_info/neuron_info.rs b/pallets/subtensor/src/rpc_info/neuron_info.rs
index b5c09ece0..781f9c5a0 100644
--- a/pallets/subtensor/src/rpc_info/neuron_info.rs
+++ b/pallets/subtensor/src/rpc_info/neuron_info.rs
@@ -28,6 +28,12 @@ pub struct NeuronInfo<T: Config> {
     pruning_score: Compact<u16>,
 }
 
+impl<T: Config> NeuronInfo<T> {
+    pub fn stake(&self) -> &[(T::AccountId, Compact<u64>)] {
+        &self.stake
+    }
+}
+
 #[freeze_struct("c21f0f4f22bcb2a1")]
 #[derive(Decode, Encode, PartialEq, Eq, Clone, Debug)]
 pub struct NeuronInfoLite<T: Config> {
@@ -52,6 +58,12 @@ pub struct NeuronInfoLite<T: Config> {
     pruning_score: Compact<u16>,
 }
 
+impl<T: Config> NeuronInfoLite<T> {
+    pub fn stake(&self) -> &[(T::AccountId, Compact<u64>)] {
+        &self.stake
+    }
+}
+
 impl<T: Config> Pallet<T> {
     pub fn get_neurons(netuid: u16) -> Vec<NeuronInfo<T>> {
         if !Self::if_subnet_exist(netuid) {
diff --git a/pallets/subtensor/src/rpc_info/stake_info.rs b/pallets/subtensor/src/rpc_info/stake_info.rs
index f5c87744d..32b47d15a 100644
--- a/pallets/subtensor/src/rpc_info/stake_info.rs
+++ b/pallets/subtensor/src/rpc_info/stake_info.rs
@@ -18,6 +18,12 @@ pub struct StakeInfo<T: Config> {
     is_registered: bool,
 }
 
+impl<T: Config> StakeInfo<T> {
+    pub fn stake(&self) -> Compact<u64> {
+        self.stake
+    }
+}
+
 impl<T: Config> Pallet<T> {
     fn _get_stake_info_for_coldkeys(
         coldkeys: Vec<T::AccountId>,
diff --git a/pallets/subtensor/tests/staking.rs b/pallets/subtensor/tests/staking.rs
index 4cd2af3b6..03cf79551 100644
--- a/pallets/subtensor/tests/staking.rs
+++ b/pallets/subtensor/tests/staking.rs
@@ -1,6 +1,7 @@
 #![allow(clippy::unwrap_used)]
 #![allow(clippy::arithmetic_side_effects)]
 
+use codec::Encode;
 use frame_support::{assert_err, assert_noop, assert_ok, traits::Currency};
 use frame_system::Config;
 mod mock;
@@ -2490,3 +2491,27 @@ fn test_anneal_global_weight() {
         );
     });
 }
+
+//  SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test staking -- test_stake_value_is_alpha --exact --nocapture
+#[test]
+fn test_stake_value_is_alpha() {
+    new_test_ext(1).execute_with(|| {
+        let hotkey_id = U256::from(5445);
+        let coldkey_id = U256::from(5443433);
+        let amount: u64 = 10000;
+        let netuid: u16 = 1;
+        let tempo: u16 = 13;
+        let start_nonce: u64 = 0;
+
+        add_network(netuid, tempo, 0);
+
+        register_ok_neuron(netuid, hotkey_id, coldkey_id, start_nonce);
+
+        SubtensorModule::stake_into_subnet(&hotkey_id, &coldkey_id, netuid, amount);
+
+        let stake_info = SubtensorModule::get_stake_info_for_coldkey(Encode::encode(&coldkey_id));
+        let neuron_info = SubtensorModule::get_neuron_lite(netuid, 0);
+
+        assert_eq!(stake_info[0].stake(), neuron_info.unwrap().stake()[0].1);
+    });
+}