From 3ce9f34d9396bfe43bfeabe3c6c20f8a168e5382 Mon Sep 17 00:00:00 2001
From: David Craven <david@craven.ch>
Date: Wed, 12 Feb 2025 10:29:32 +0100
Subject: [PATCH] Refactor faucet.

---
 gmp/evm/src/lib.rs    | 11 +----------
 gmp/grpc/src/lib.rs   |  4 ++--
 gmp/grpc/src/main.rs  |  7 +++++--
 gmp/grpc/src/proto.rs |  4 +++-
 gmp/rust/src/lib.rs   |  6 ++----
 primitives/src/gmp.rs |  2 +-
 tc-cli/src/config.rs  |  1 +
 tc-cli/src/lib.rs     |  4 +++-
 8 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/gmp/evm/src/lib.rs b/gmp/evm/src/lib.rs
index 32ac21845..5ea2308e1 100644
--- a/gmp/evm/src/lib.rs
+++ b/gmp/evm/src/lib.rs
@@ -426,16 +426,7 @@ impl IChain for Connector {
 		(config.currency_decimals, config.currency_symbol)
 	}
 	/// Uses a faucet to fund the account when possible.
-	async fn faucet(&self) -> Result<()> {
-		let balance = match self.network_id() {
-			6 => 10u128.pow(25), // astar
-			2 => 10u128.pow(29), // ethereum
-			3 => 10u128.pow(29), // ethereum
-			network_id => {
-				tracing::info!("network {network_id} doesn't support faucet");
-				return Ok(());
-			},
-		};
+	async fn faucet(&self, balance: u128) -> Result<()> {
 		self.wallet.faucet(balance, None).await?;
 		Ok(())
 	}
diff --git a/gmp/grpc/src/lib.rs b/gmp/grpc/src/lib.rs
index 65a668d6a..924103a89 100644
--- a/gmp/grpc/src/lib.rs
+++ b/gmp/grpc/src/lib.rs
@@ -98,8 +98,8 @@ impl IChain for Connector {
 		self.address
 	}
 	/// Uses a faucet to fund the account when possible.
-	async fn faucet(&self) -> Result<()> {
-		let request = Request::new(proto::FaucetRequest {});
+	async fn faucet(&self, balance: u128) -> Result<()> {
+		let request = Request::new(proto::FaucetRequest { balance });
 		self.client.lock().await.faucet(request).await?;
 		Ok(())
 	}
diff --git a/gmp/grpc/src/main.rs b/gmp/grpc/src/main.rs
index b3fd032d5..ba1875122 100644
--- a/gmp/grpc/src/main.rs
+++ b/gmp/grpc/src/main.rs
@@ -54,8 +54,11 @@ impl Gmp for ConnectorWrapper {
 		&self,
 		request: Request<proto::FaucetRequest>,
 	) -> GmpResult<proto::FaucetResponse> {
-		let (connector, _) = self.connector(request)?;
-		connector.faucet().await.map_err(|err| Status::unknown(err.to_string()))?;
+		let (connector, msg) = self.connector(request)?;
+		connector
+			.faucet(msg.balance)
+			.await
+			.map_err(|err| Status::unknown(err.to_string()))?;
 		Ok(Response::new(proto::FaucetResponse {}))
 	}
 
diff --git a/gmp/grpc/src/proto.rs b/gmp/grpc/src/proto.rs
index 350d8dae4..81d783b89 100644
--- a/gmp/grpc/src/proto.rs
+++ b/gmp/grpc/src/proto.rs
@@ -6,7 +6,9 @@ use time_primitives::{
 };
 
 #[derive(Serialize, Deserialize)]
-pub struct FaucetRequest {}
+pub struct FaucetRequest {
+	pub balance: u128,
+}
 
 #[derive(Serialize, Deserialize)]
 pub struct FaucetResponse {}
diff --git a/gmp/rust/src/lib.rs b/gmp/rust/src/lib.rs
index d5d8c9926..d9a899da9 100644
--- a/gmp/rust/src/lib.rs
+++ b/gmp/rust/src/lib.rs
@@ -23,7 +23,6 @@ use time_primitives::{
 
 const BLOCK_TIME: u64 = 1;
 const FINALIZATION_TIME: u64 = 2;
-const FAUCET: u128 = 1_000_000_000;
 
 const BLOCKS: TableDefinition<u64, u64> = TableDefinition::new("blocks");
 const BALANCE: TableDefinition<Address, u128> = TableDefinition::new("balance");
@@ -193,12 +192,11 @@ impl IChain for Connector {
 		currency()
 	}
 
-	async fn faucet(&self) -> Result<()> {
+	async fn faucet(&self, balance: u128) -> Result<()> {
 		let tx = self.db.begin_write()?;
 		{
 			let mut t = tx.open_table(BALANCE)?;
-			let balance = read_balance(&t, self.address)?;
-			t.insert(self.address, balance + FAUCET)?;
+			t.insert(self.address, balance)?;
 		}
 		tx.commit()?;
 		Ok(())
diff --git a/primitives/src/gmp.rs b/primitives/src/gmp.rs
index d8e9ba010..c42210a2b 100644
--- a/primitives/src/gmp.rs
+++ b/primitives/src/gmp.rs
@@ -354,7 +354,7 @@ pub trait IChain: Send + Sync + 'static {
 	/// Human readable connector account identifier.
 	fn address(&self) -> Address;
 	/// Uses a faucet to fund the account when possible.
-	async fn faucet(&self) -> Result<()>;
+	async fn faucet(&self, balance: u128) -> Result<()>;
 	/// Transfers an amount to an account.
 	async fn transfer(&self, address: Address, amount: u128) -> Result<()>;
 	/// Queries the account balance.
diff --git a/tc-cli/src/config.rs b/tc-cli/src/config.rs
index cefafc845..c71dfbd8a 100644
--- a/tc-cli/src/config.rs
+++ b/tc-cli/src/config.rs
@@ -135,6 +135,7 @@ pub struct NetworkConfig {
 	pub route_base_fee: u128,
 	pub shard_size: u16,
 	pub shard_threshold: u16,
+	pub faucet: Option<u128>,
 }
 
 #[cfg(test)]
diff --git a/tc-cli/src/lib.rs b/tc-cli/src/lib.rs
index e96ac31e1..c17df8849 100644
--- a/tc-cli/src/lib.rs
+++ b/tc-cli/src/lib.rs
@@ -215,7 +215,9 @@ impl Tc {
 	}
 
 	pub async fn faucet(&self, network: NetworkId) -> Result<()> {
-		self.connector(network)?.faucet().await
+		let config = self.config.network(network)?;
+		let faucet = config.faucet.context("faucet not supported, please prefund account")?;
+		self.connector(network)?.faucet(faucet).await
 	}
 
 	pub async fn balance(&self, network: Option<NetworkId>, address: Address) -> Result<u128> {