From 9cca9a12c8ad3ab9c57432365e6e8044855c5b77 Mon Sep 17 00:00:00 2001 From: shinorh247 <304154597+shinorh247@users.noreply.github.com> Date: Fri, 2 Oct 2026 09:12:25 +0000 Subject: [PATCH] Report keys staked via stake_keys in get_total_staked The two staking flows maintain separate aggregates: stake_keys_locked and friends maintain StakingRewardsState::total_staked next to the rewards pool, while stake_keys/unstake_keys maintain DataKey::TotalStaked. The view read only the former, so a creator whose keys were staked through the simple flow reported a total of 0. Sum both aggregates so the view matches what the contract actually tracks. Neither writer changes, so the locked flow and the rewards math that shares StakingRewardsState are unaffected. Co-Authored-By: Codebuff --- creator-keys/src/lib.rs | 16 ++- ...total_staked_reports_both_staking_flows.rs | 107 ++++++++++++++++++ 2 files changed, 121 insertions(+), 2 deletions(-) create mode 100644 creator-keys/tests/get_total_staked_reports_both_staking_flows.rs diff --git a/creator-keys/src/lib.rs b/creator-keys/src/lib.rs index c8497787..2db6c0be 100644 --- a/creator-keys/src/lib.rs +++ b/creator-keys/src/lib.rs @@ -11639,14 +11639,26 @@ impl CreatorKeysContract { /// Read-only view: returns the total number of keys currently staked for /// `creator` across all holders. + /// + /// Two staking flows maintain separate aggregates and are summed here so + /// the view reflects both: + /// - `stake_keys_locked` / `early_unstake` / `claim_stake_reward` maintain + /// `StakingRewardsState::total_staked` alongside the rewards pool. + /// - `stake_keys` / `unstake_keys` maintain `DataKey::TotalStaked`. + /// + /// Reading only the former made this view report `0` for keys staked + /// through the simple flow. pub fn get_total_staked(env: Env, creator: Address) -> u32 { - env.storage() + let locked_total: u32 = env + .storage() .persistent() .get::(&constants::storage::staking_rewards_pool( &creator, )) .map(|state| state.total_staked) - .unwrap_or(0) + .unwrap_or(0); + let simple_total: u32 = read_total_staked(&env, &creator); + locked_total.saturating_add(simple_total) } // ========================================================================= diff --git a/creator-keys/tests/get_total_staked_reports_both_staking_flows.rs b/creator-keys/tests/get_total_staked_reports_both_staking_flows.rs new file mode 100644 index 00000000..c089b303 --- /dev/null +++ b/creator-keys/tests/get_total_staked_reports_both_staking_flows.rs @@ -0,0 +1,107 @@ +//! Regression test: `get_total_staked` must report keys staked through the +//! simple `stake_keys` flow, not only the locked `stake_keys_locked` flow. +//! +//! The two flows maintain separate aggregates: +//! - `stake_keys_locked` / `early_unstake` / `claim_stake_reward` maintain +//! `StakingRewardsState::total_staked` (alongside the rewards pool). +//! - `stake_keys` / `unstake_keys` maintain `DataKey::TotalStaked`. +//! +//! The view previously read only the former, so it reported `0` for keys +//! staked through the simple flow. + +mod contract_test_env; + +use contract_test_env::{ + register_creator_keys, register_test_creator, set_pricing_and_fees, test_env_with_auths, +}; +use soroban_sdk::{testutils::Address as _, Address}; + +const KEY_PRICE: i128 = 1_000; +const CREATOR_BPS: u32 = 9_000; +const PROTOCOL_BPS: u32 = 1_000; +const LOCK_LEDGERS: u32 = 100; + +/// Buys `keys` keys for `holder`, one buy at a time. +fn buy_keys( + client: &creator_keys::CreatorKeysContractClient<'_>, + creator: &Address, + holder: &Address, + keys: u32, +) { + for _ in 0..keys { + client.buy_key(creator, holder, &KEY_PRICE, &None); + } +} + +#[test] +fn simple_stake_flow_is_visible_to_get_total_staked() { + let env = test_env_with_auths(); + let (client, _) = register_creator_keys(&env); + set_pricing_and_fees(&env, &client, KEY_PRICE, CREATOR_BPS, PROTOCOL_BPS); + let creator = register_test_creator(&env, &client, "alice"); + let holder = Address::generate(&env); + + assert_eq!(client.get_total_staked(&creator), 0); + + buy_keys(&client, &creator, &holder, 1); + client.stake_keys(&creator, &holder, &1u32); + + assert_eq!( + client.get_staked_balance(&creator, &holder), + 1, + "the simple flow records the per-wallet staked balance", + ); + assert_eq!( + client.get_total_staked(&creator), + 1, + "get_total_staked must reflect keys staked via stake_keys", + ); + + // Unstaking the last key must bring the aggregate back to zero. + client.unstake_keys(&creator, &holder, &1u32); + assert_eq!(client.get_total_staked(&creator), 0); +} + +#[test] +fn locked_stake_flow_remains_visible_to_get_total_staked() { + let env = test_env_with_auths(); + let (client, _) = register_creator_keys(&env); + set_pricing_and_fees(&env, &client, KEY_PRICE, CREATOR_BPS, PROTOCOL_BPS); + let creator = register_test_creator(&env, &client, "bob"); + let holder = Address::generate(&env); + + buy_keys(&client, &creator, &holder, 4); + client.stake_keys_locked(&creator, &holder, &3u32, &LOCK_LEDGERS); + + assert_eq!( + client.get_total_staked(&creator), + 3, + "get_total_staked must still reflect keys staked via stake_keys_locked", + ); +} + +#[test] +fn both_staking_flows_are_summed() { + let env = test_env_with_auths(); + let (client, _) = register_creator_keys(&env); + set_pricing_and_fees(&env, &client, KEY_PRICE, CREATOR_BPS, PROTOCOL_BPS); + let creator = register_test_creator(&env, &client, "carol"); + let holder_a = Address::generate(&env); + let holder_b = Address::generate(&env); + + // holder_a stakes 2 via the simple flow. + buy_keys(&client, &creator, &holder_a, 2); + client.stake_keys(&creator, &holder_a, &2u32); + + // holder_b stakes 3 via the locked flow. + buy_keys(&client, &creator, &holder_b, 3); + client.stake_keys_locked(&creator, &holder_b, &3u32, &LOCK_LEDGERS); + + assert_eq!(client.get_staked_balance(&creator, &holder_a), 2); + assert_eq!(client.get_staked_balance(&creator, &holder_b), 3); + assert_eq!( + client.get_total_staked(&creator), + 5, + "both flows must contribute to the same aggregate", + ); +}