feat: add chain.get_balance() method (#2520)

Author: antazoey

src/ape/managers/chain.py

@@ -22,6 +22,7 @@
     APINotImplementedError,
     BlockNotFoundError,
     ChainError,
+    ConversionError,
     ProviderNotConnectedError,
     QueryEngineError,
     TransactionNotFoundError,
@@ -937,7 +938,39 @@ def mine(
             self.pending_timestamp += deltatime
         self.provider.mine(num_blocks)
 
-    def set_balance(self, account: Union[BaseAddress, AddressType], amount: Union[int, str]):
+    def get_balance(
+        self, address: Union[BaseAddress, AddressType, str], block_id: Optional["BlockID"] = None
+    ) -> int:
+        """
+        Get the balance of the given address. If ``ape-ens`` is installed,
+        you can pass ENS names.
+
+        Args:
+            address (BaseAddress, AddressType | str): An address, ENS, or account/contract.
+            block_id (:class:`~ape.types.BlockID` | None): The block ID. Defaults to latest.
+
+        Returns:
+            int: The account balance.
+        """
+        if (isinstance(address, str) and not address.startswith("0x")) or not isinstance(
+            address, str
+        ):
+            try:
+                address = self.conversion_manager.convert(address, AddressType)
+            except ConversionError:
+                # Try to get the balance anyway; maybe the provider can handle it.
+                address = address
+
+        return self.provider.get_balance(address, block_id=block_id)
+
+    def set_balance(self, account: Union[BaseAddress, AddressType, str], amount: Union[int, str]):
+        """
+        Set an account balance, only works on development chains.
+
+        Args:
+            account (BaseAddress, AddressType | str): The account.
+            amount (int | str): The new balance.
+        """
         if isinstance(account, BaseAddress):
             account = account.address
 

tests/functional/test_chain.py

@@ -173,6 +173,11 @@ def test_set_pending_timestamp_failure(chain):
         )
 
 
+def test_get_balance(chain, owner):
+    assert chain.get_balance(owner) == owner.balance
+    assert chain.get_balance(owner.address) == owner.balance
+
+
 def test_set_balance(chain, owner):
     with pytest.raises(APINotImplementedError):
         chain.set_balance(owner, "1000 ETH")