allow specify balances for each users and each tokens (#66)

Author: ermalkaleci

currencies/src/mock.rs

@@ -97,33 +97,33 @@ pub const BOB: AccountId = 2;
 pub const EVA: AccountId = 5;
 
 pub struct ExtBuilder {
-	currency_ids: Vec<CurrencyId>,
-	endowed_accounts: Vec<AccountId>,
-	initial_balance: Balance,
+	endowed_accounts: Vec<(AccountId, CurrencyId, Balance)>,
 	// whether the configs are for `pallet_balances` or not
 	is_for_pallet_balances: bool,
 }
 
 impl Default for ExtBuilder {
 	fn default() -> Self {
 		Self {
-			currency_ids: vec![NATIVE_CURRENCY_ID, X_TOKEN_ID],
-			endowed_accounts: vec![0],
-			initial_balance: 0,
+			endowed_accounts: vec![],
 			is_for_pallet_balances: false,
 		}
 	}
 }
 
 impl ExtBuilder {
-	pub fn balances(mut self, account_ids: Vec<AccountId>, initial_balance: Balance) -> Self {
-		self.endowed_accounts = account_ids;
-		self.initial_balance = initial_balance;
+	pub fn balances(mut self, endowed_accounts: Vec<(AccountId, CurrencyId, Balance)>) -> Self {
+		self.endowed_accounts = endowed_accounts;
 		self
 	}
 
 	pub fn one_hundred_for_alice_n_bob(self) -> Self {
-		self.balances(vec![ALICE, BOB], 100)
+		self.balances(vec![
+			(ALICE, NATIVE_CURRENCY_ID, 100),
+			(BOB, NATIVE_CURRENCY_ID, 100),
+			(ALICE, X_TOKEN_ID, 100),
+			(BOB, X_TOKEN_ID, 100),
+		])
 	}
 
 	pub fn make_for_pallet_balances(mut self) -> Self {
@@ -140,17 +140,16 @@ impl ExtBuilder {
 			pallet_balances::GenesisConfig::<Runtime> {
 				balances: self
 					.endowed_accounts
-					.iter()
-					.map(|acc| (acc.clone(), self.initial_balance))
-					.collect(),
+					.into_iter()
+					.filter(|(_, currency_id, _)| *currency_id == X_TOKEN_ID)
+					.map(|(account_id, _, initial_balance)| (account_id, initial_balance))
+					.collect::<Vec<_>>(),
 				vesting: vec![],
 			}
 			.assimilate_storage(&mut t)
 			.unwrap();
 		} else {
 			tokens::GenesisConfig::<Runtime> {
-				tokens: self.currency_ids,
-				initial_balance: self.initial_balance,
 				endowed_accounts: self.endowed_accounts,
 			}
 			.assimilate_storage(&mut t)

tokens/src/lib.rs

@@ -1,7 +1,10 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use frame_support::{decl_error, decl_event, decl_module, decl_storage, ensure, Parameter};
-use rstd::convert::{TryFrom, TryInto};
+use rstd::{
+	collections::btree_map::BTreeMap,
+	convert::{TryFrom, TryInto},
+};
 use sp_runtime::{
 	traits::{CheckedAdd, CheckedSub, MaybeSerializeDeserialize, Member, SimpleArithmetic, StaticLookup},
 	DispatchResult,
@@ -31,30 +34,38 @@ pub trait Trait: frame_system::Trait {
 		+ Default
 		+ Copy
 		+ MaybeSerializeDeserialize;
-	type CurrencyId: Parameter + Member + Copy + MaybeSerializeDeserialize;
+	type CurrencyId: Parameter + Member + Copy + MaybeSerializeDeserialize + Ord;
 }
 
 decl_storage! {
 	trait Store for Module<T: Trait> as Tokens {
 		/// The total issuance of a token type.
 		pub TotalIssuance get(fn total_issuance) build(|config: &GenesisConfig<T>| {
-			let issuance = config.initial_balance * (config.endowed_accounts.len() as u32).into();
-			config.tokens.iter().map(|id| (id.clone(), issuance)).collect::<Vec<_>>()
+			config
+				.endowed_accounts
+				.iter()
+				.map(|(_, currency_id, initial_balance)| (currency_id, initial_balance))
+				.fold(BTreeMap::<T::CurrencyId, T::Balance>::new(), |mut acc, (currency_id, initial_balance)| {
+					if let Some(issuance) = acc.get_mut(currency_id) {
+						*issuance = issuance.checked_add(initial_balance).expect("total issuance cannot overflow when building genesis");
+					} else {
+						acc.insert(*currency_id, *initial_balance);
+					}
+					acc
+				})
+				.into_iter()
+				.collect::<Vec<_>>()
 		}): map T::CurrencyId => T::Balance;
 
 		/// The balance of a token type under an account.
 		pub Balance get(fn balance): double_map T::CurrencyId, blake2_256(T::AccountId) => T::Balance;
 	}
 	add_extra_genesis {
-		config(tokens): Vec<T::CurrencyId>;
-		config(initial_balance): T::Balance;
-		config(endowed_accounts): Vec<T::AccountId>;
+		config(endowed_accounts): Vec<(T::AccountId, T::CurrencyId, T::Balance)>;
 
 		build(|config: &GenesisConfig<T>| {
-			config.tokens.iter().for_each(|currency_id| {
-				config.endowed_accounts.iter().for_each(|account_id| {
-					<Balance<T>>::insert(currency_id, account_id, &config.initial_balance);
-				})
+			config.endowed_accounts.iter().for_each(|(account_id, currency_id, initial_balance)| {
+				<Balance<T>>::insert(currency_id, account_id, initial_balance);
 			})
 		})
 	}

tokens/src/mock.rs

@@ -36,9 +36,9 @@ parameter_types! {
 type AccountId = u64;
 impl frame_system::Trait for Runtime {
 	type Origin = Origin;
+	type Call = ();
 	type Index = u64;
 	type BlockNumber = u64;
-	type Call = ();
 	type Hash = H256;
 	type Hashing = ::sp_runtime::traits::BlakeTwo256;
 	type AccountId = AccountId;
@@ -70,30 +70,25 @@ pub const ALICE: AccountId = 1;
 pub const BOB: AccountId = 2;
 
 pub struct ExtBuilder {
-	currency_id: CurrencyId,
-	endowed_accounts: Vec<AccountId>,
-	initial_balance: Balance,
+	endowed_accounts: Vec<(AccountId, CurrencyId, Balance)>,
 }
 
 impl Default for ExtBuilder {
 	fn default() -> Self {
 		Self {
-			currency_id: TEST_TOKEN_ID,
-			endowed_accounts: vec![0],
-			initial_balance: 0,
+			endowed_accounts: vec![],
 		}
 	}
 }
 
 impl ExtBuilder {
-	pub fn balances(mut self, account_ids: Vec<AccountId>, initial_balance: Balance) -> Self {
-		self.endowed_accounts = account_ids;
-		self.initial_balance = initial_balance;
+	pub fn balances(mut self, endowed_accounts: Vec<(AccountId, CurrencyId, Balance)>) -> Self {
+		self.endowed_accounts = endowed_accounts;
 		self
 	}
 
 	pub fn one_hundred_for_alice_n_bob(self) -> Self {
-		self.balances(vec![ALICE, BOB], 100)
+		self.balances(vec![(ALICE, TEST_TOKEN_ID, 100), (BOB, TEST_TOKEN_ID, 100)])
 	}
 
 	pub fn build(self) -> runtime_io::TestExternalities {
@@ -102,8 +97,6 @@ impl ExtBuilder {
 			.unwrap();
 
 		GenesisConfig::<Runtime> {
-			tokens: vec![self.currency_id],
-			initial_balance: self.initial_balance,
 			endowed_accounts: self.endowed_accounts,
 		}
 		.assimilate_storage(&mut t)