From 3163d5f4e62728a4ff5aa312b4302ec38bec6d4c Mon Sep 17 00:00:00 2001
From: runcomet <runcomet@protonmail.com>
Date: Tue, 12 Nov 2024 18:48:08 -0800
Subject: [PATCH] configure ensure_ti_valid awareness of dev_accounts

---
 .../frame/balances/src/tests/general_tests.rs | 13 -----
 substrate/frame/balances/src/tests/mod.rs     | 53 +++++++++++++++----
 2 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/substrate/frame/balances/src/tests/general_tests.rs b/substrate/frame/balances/src/tests/general_tests.rs
index d062d8862914..a855fae5616a 100644
--- a/substrate/frame/balances/src/tests/general_tests.rs
+++ b/substrate/frame/balances/src/tests/general_tests.rs
@@ -141,16 +141,3 @@ fn try_state_works() {
 			.contains("Found `Freeze` with too many elements"));
 	});
 }
-
-#[cfg(feature = "runtime-benchmarks")]
-#[test]
-fn dev_accounts_populated() {
-    ExtBuilder::default().build_and_execute_with(|| {
-        UseSystem::set(true); // copmment this out and uncomment this from time to time to check stuff.
-        // Print total issuance for debugging
-        let ti = TotalIssuance::<Test>::get();
-        println!("Total Issuance after genesis: {}", ti);
-		println!("Why are you printing twice?");
-        ensure_ti_valid();
-    });
-}
diff --git a/substrate/frame/balances/src/tests/mod.rs b/substrate/frame/balances/src/tests/mod.rs
index a621b8c21f9f..8b953b957291 100644
--- a/substrate/frame/balances/src/tests/mod.rs
+++ b/substrate/frame/balances/src/tests/mod.rs
@@ -43,6 +43,11 @@ use sp_runtime::{
 };
 use std::collections::BTreeSet;
 
+#[cfg(feature = "runtime-benchmarks")]
+use sp_core::{sr25519::Pair as SrPair, Pair};
+
+use frame_support::traits::BuildGenesisConfig;
+
 mod currency_tests;
 mod dispatchable_tests;
 mod fungible_conformance_tests;
@@ -281,19 +286,45 @@ pub fn info_from_weight(w: Weight) -> DispatchInfo {
 
 /// Check that the total-issuance matches the sum of all accounts' total balances.
 pub fn ensure_ti_valid() {
-	let mut sum = 0;
+    let mut sum = 0;
 
-	for acc in frame_system::Account::<Test>::iter_keys() {
-		if UseSystem::get() {
-			let data = frame_system::Pallet::<Test>::account(acc);
-			sum += data.data.total();
-		} else {
-			let data = crate::Account::<Test>::get(acc);
-			sum += data.total();
-		}
-	}
+    // Fetch the dev accounts from the GenesisConfig.
+    #[cfg(feature = "runtime-benchmarks")]
+    let dev_accounts = (10, 100, "//Sender/{}".to_string()); // You can customize this as needed
+    #[cfg(feature = "runtime-benchmarks")]
+    let (num_accounts, balance, ref derivation) = dev_accounts;
+
+    // Generate the dev account public keys.
+    #[cfg(feature = "runtime-benchmarks")]
+    let dev_account_ids: Vec<_> = (0..num_accounts)
+        .map(|index| {
+            let derivation_string = derivation.replace("{}", &index.to_string());
+            let pair: SrPair = Pair::from_string(&derivation_string, None).expect("Invalid derivation string");
+            <crate::tests::Test as frame_system::Config>::AccountId::decode(&mut &pair.public().encode()[..]).unwrap()
+        })
+        .collect();
+
+    // Iterate over all account keys (i.e., the account IDs).
+    for acc in frame_system::Account::<Test>::iter_keys() {
+        // Skip dev accounts by checking if the account is in the dev_account_ids list.
+		// This also proves dev_accounts exists in storage.
+        #[cfg(feature = "runtime-benchmarks")]
+        if dev_account_ids.contains(&acc) {
+            continue;
+        }
+
+        // Check if we are using the system pallet or some other custom storage for accounts.
+        if UseSystem::get() {
+            let data = frame_system::Pallet::<Test>::account(acc);
+            sum += data.data.total();
+        } else {
+            let data = crate::Account::<Test>::get(acc);
+            sum += data.total();
+        }
+    }
 
-	assert_eq!(TotalIssuance::<Test>::get(), sum, "Total Issuance wrong");
+    // Ensure the total issuance matches the sum of the account balances
+    assert_eq!(TotalIssuance::<Test>::get(), sum, "Total Issuance is incorrect");
 }
 
 #[test]