feat: automate faucet claims for test agents (#197)

python/examples/06-send-tokens/main.py

@@ -19,8 +19,8 @@ class TransactionInfo(Model):
 alice = Agent(name="alice", seed="alice secret phrase")
 bob = Agent(name="bob", seed="bob secret phrase")
 
-fund_agent_if_low(alice.wallet.address())
-fund_agent_if_low(bob.wallet.address())
+
+fund_agent_if_low(bob.wallet.address(), min_balance=AMOUNT)
 
 
 @alice.on_interval(period=10.0)

python/examples/08-local-network-interaction/agent1.py

@@ -1,4 +1,3 @@
-from uagents.setup import fund_agent_if_low
 from uagents import Agent, Context, Model
 
 
@@ -16,8 +15,6 @@ class Message(Model):
  endpoint=["http://127.0.0.1:8001/submit"],
 )
 
-fund_agent_if_low(bob.wallet.address())
-
 
 @bob.on_message(model=Message)
 async def message_handler(ctx: Context, sender: str, msg: Message):

python/examples/08-local-network-interaction/agent2.py

@@ -1,4 +1,3 @@
-from uagents.setup import fund_agent_if_low
 from uagents import Agent, Context, Model
 
 
@@ -17,8 +16,6 @@ class Message(Model):
  endpoint=["http://127.0.0.1:8000/submit"],
 )
 
-fund_agent_if_low(alice.wallet.address())
-
 
 @alice.on_interval(period=2.0)
 async def send_message(ctx: Context):

python/examples/09-booking-protocol-demo/restaurant.py

@@ -3,7 +3,6 @@
 
 from uagents import Agent
 from uagents.contrib.protocols.protocol_query import proto_query
-from uagents.setup import fund_agent_if_low
 
 
 restaurant = Agent(
@@ -15,7 +14,6 @@
  },
 )
 
-fund_agent_if_low(restaurant.wallet.address())
 
 # build the restaurant agent from stock protocols and publish their details
 restaurant.include(query_proto, publish_manifest=True)

python/examples/09-booking-protocol-demo/user.py

@@ -5,7 +5,6 @@
 )
 
 from uagents import Agent, Context
-from uagents.setup import fund_agent_if_low
 
 
 RESTAURANT_ADDRESS = (
@@ -21,7 +20,6 @@
  },
 )
 
-fund_agent_if_low(user.wallet.address())
 
 table_query = QueryTableRequest(
  guests=3,

python/examples/10-cleaning-demo/cleaner.py

@@ -7,7 +7,6 @@
 from protocols.cleaning.models import Availability, Provider, Service, ServiceType
 
 from uagents import Agent, Context
-from uagents.setup import fund_agent_if_low
 
 
 cleaner = Agent(
@@ -19,8 +18,6 @@
  },
 )
 
-fund_agent_if_low(cleaner.wallet.address())
-
 
 # build the cleaning service agent from the cleaning protocol
 cleaner.include(cleaning_proto)

python/examples/10-cleaning-demo/user.py

@@ -9,7 +9,6 @@
 )
 from protocols.cleaning.models import ServiceType
 from uagents import Agent, Context
-from uagents.setup import fund_agent_if_low
 
 
 CLEANER_ADDRESS = (
@@ -25,7 +24,6 @@
  },
 )
 
-fund_agent_if_low(user.wallet.address())
 
 request = ServiceRequest(
  user=user.name,

python/examples/11-mailbox-agents/alice.py

@@ -1,5 +1,4 @@
 from uagents import Agent, Context, Model
-from uagents.setup import fund_agent_if_low
 
 
 class Message(Model):
@@ -22,8 +21,6 @@ class Message(Model):
  mailbox=f"{API_KEY}@https://agentverse.ai",
 )
 
-fund_agent_if_low(agent.wallet.address())
-
 
 @agent.on_message(model=Message, replies={Message})
 async def handle_message(ctx: Context, sender: str, msg: Message):

python/examples/11-mailbox-agents/bob.py

@@ -1,5 +1,4 @@
 from uagents import Agent, Context, Model
-from uagents.setup import fund_agent_if_low
 
 
 class Message(Model):
@@ -25,8 +24,6 @@ class Message(Model):
  mailbox=f"{API_KEY}@https://agentverse.ai",
 )
 
-fund_agent_if_low(agent.wallet.address())
-
 
 @agent.on_interval(period=2.0)
 async def send_message(ctx: Context):

python/examples/12-remote-agents/agent1.py

@@ -10,7 +10,6 @@
 from pyngrok import conf, ngrok
 
 from uagents import Agent, Context, Model
-from uagents.setup import fund_agent_if_low
 
 
 class Message(Model):
@@ -29,8 +28,6 @@ class Message(Model):
  endpoint=[f"{http_tunnel.public_url}/submit"],
 )
 
-fund_agent_if_low(alice.wallet.address())
-
 
 @alice.on_message(model=Message)
 async def act_on_message(ctx: Context, sender: str, msg: Message):

python/examples/12-remote-agents/agent2.py

@@ -2,7 +2,6 @@
 Counterpart for the remote connection example in agent1.py.
 """
 from uagents import Agent, Context, Model
-from uagents.setup import fund_agent_if_low
 
 
 class Message(Model):
@@ -16,8 +15,6 @@ class Message(Model):
  endpoint=["http://127.0.0.1:8001/submit"],
 )
 
-fund_agent_if_low(bob.wallet.address())
-
 ALICE_ADDRESS = "agent1qv2l7qzcd2g2rcv2p93tqflrcaq5dk7c2xc7fcnfq3s37zgkhxjmq5mfyvz"
 
 

python/examples/13-agent-name-service/agent1.py

@@ -1,7 +1,6 @@
 from cosmpy.aerial.wallet import LocalWallet
 
-from uagents.network import get_name_service_contract
-from uagents.setup import fund_agent_if_low
+from uagents.network import get_name_service_contract, get_faucet
 from uagents import Agent, Context, Model
 
 
@@ -22,11 +21,10 @@ class Message(Model):
 
 my_wallet = LocalWallet.from_unsafe_seed("registration test wallet")
 name_service_contract = get_name_service_contract(test=True)
+faucet = get_faucet()
 DOMAIN = "agent"
 
-
-for wallet in [my_wallet, bob.wallet]:
- fund_agent_if_low(wallet.address())
+faucet.get_wealth(my_wallet.address())
 
 
 @bob.on_event("startup")

python/examples/13-agent-name-service/agent2.py

@@ -1,4 +1,3 @@
-from uagents.setup import fund_agent_if_low
 from uagents import Agent, Context, Model
 
 
@@ -13,8 +12,6 @@ class Message(Model):
  endpoint=["http://localhost:8000/submit"],
 )
 
-fund_agent_if_low(alice.wallet.address())
-
 
 @alice.on_interval(period=5)
 async def alice_interval_handler(ctx: Context):

python/examples/14-broadcast/main.py

@@ -1,5 +1,4 @@
 from uagents import Agent, Bureau, Context, Model, Protocol
-from uagents.setup import fund_agent_if_low
 
 # create agents
 # alice and bob will support the protocol
@@ -8,10 +7,6 @@
 bob = Agent(name="bob", seed="bob recovery phrase")
 charles = Agent(name="charles", seed="charles recovery phrase")
 
-fund_agent_if_low(alice.wallet.address())
-fund_agent_if_low(bob.wallet.address())
-fund_agent_if_low(charles.wallet.address())
-
 
 class BroadcastExampleRequest(Model):
  pass

python/src/uagents/agent.py

@@ -28,9 +28,12 @@
 from uagents.network import (
  get_ledger,
  get_almanac_contract,
+ add_testnet_funds,
+ InsufficientFundsError,
 )
 from uagents.mailbox import MailboxClient
 from uagents.config import (
+ AVERAGE_BLOCK_INTERVAL,
  REGISTRATION_FEE,
  REGISTRATION_UPDATE_INTERVAL_SECONDS,
  LEDGER_PREFIX,
@@ -512,10 +515,18 @@ async def register(self):
  ):
  if self.balance < REGISTRATION_FEE:
  self._logger.warning(
- f"I do not have enough funds to register on Almanac contract\
- \nFund using wallet address: {self.wallet.address()}"
+ "I do not have enough funds to register on Almanac contract"
  )
- return
+ if self._test:
+ add_testnet_funds(str(self.wallet.address()))
+ self._logger.info(
+ f"Adding testnet funds to {self.wallet.address()}"
+ )
+ else:
+ self._logger.info(
+ f"Send funds to wallet address: {self.wallet.address()}"
+ )
+ raise InsufficientFundsError()
  self._logger.info("Registering on almanac contract...")
  signature = self.sign_registration()
  await self._almanac_contract.register(
@@ -541,6 +552,8 @@ async def _registration_loop(self):
  time_until_next_registration = REGISTRATION_UPDATE_INTERVAL_SECONDS
  try:
  await self.register()
+ except InsufficientFundsError:
+ time_until_next_registration = 2 * AVERAGE_BLOCK_INTERVAL
  except Exception as ex:
  self._logger.exception(f"Failed to register on almanac contract: {ex}")
  time_until_next_registration = REGISTRATION_RETRY_INTERVAL_SECONDS

python/src/uagents/config.py

@@ -31,7 +31,7 @@
 REGISTRATION_DENOM = "atestfet"
 REGISTRATION_UPDATE_INTERVAL_SECONDS = 3600
 REGISTRATION_RETRY_INTERVAL_SECONDS = 60
-AVERAGE_BLOCK_INTERVAL = 5.7
+AVERAGE_BLOCK_INTERVAL = 6
 
 AGENTVERSE_URL = "https://agentverse.ai"
 ALMANAC_API_URL = AGENTVERSE_URL + "/v1/almanac/"

python/src/uagents/network.py

@@ -39,6 +39,10 @@
 _mainnet_ledger = LedgerClient(NetworkConfig.fetchai_mainnet())
 
 
+class InsufficientFundsError(Exception):
+ """Raised when an agent has insufficient funds for a transaction."""
+
+
 def get_ledger(test: bool = True) -> LedgerClient:
  """
  Get the Ledger client.
@@ -64,6 +68,18 @@ def get_faucet() -> FaucetApi:
  return _faucet_api
 
 
+def add_testnet_funds(wallet_address: str):
+ """
+ Add testnet funds to the provided wallet address.
+
+ Args:
+ wallet_address (str): The wallet address to add funds to.
+ """
+ _faucet_api._try_create_faucet_claim( # pylint: disable=protected-access
+ wallet_address
+ )
+
+
 async def wait_for_tx_to_complete(
  tx_hash: str,
  ledger: LedgerClient,

python/src/uagents/setup.py

@@ -10,12 +10,13 @@
 LOGGER = get_logger("setup")
 
 
-def fund_agent_if_low(wallet_address: str):
+def fund_agent_if_low(wallet_address: str, min_balance: int = REGISTRATION_FEE):
  """
- Checks the agent's wallet balance and adds testnet funds if it's below the registration fee.
+ Checks the agent's wallet balance and adds testnet funds if it's below min_balance.
 
  Args:
  wallet_address (str): The wallet address of the agent.
+ min_balance (int): The minimum balance required.
 
  Returns:
  None
@@ -25,7 +26,7 @@ def fund_agent_if_low(wallet_address: str):
 
  agent_balance = ledger.query_bank_balance(Address(wallet_address))
 
- if agent_balance < REGISTRATION_FEE:
+ if agent_balance < min_balance:
  try:
  LOGGER.info("Adding testnet funds to agent...")
  faucet.get_wealth(wallet_address)