remove resolve_address functionality from safetxbuilder/tests

BalancerMaxis · Jan 10, 2025 · 0101010 · 0101010
1 parent 145d335
commit 0101010
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 40 deletions.
diff --git a/bal_tools/safe_tx_builder/safe_contract.py b/bal_tools/safe_tx_builder/safe_contract.py
@@ -18,7 +18,7 @@ def __init__(
         self.tx_builder = SafeTxBuilder()
         if not self.tx_builder._initialized:
             raise RuntimeError("SafeTxBuilder must be initialized before using SafeContract")
-        self.address = self.tx_builder._resolve_address(address)
+        self.address = address
         self.abi = self._load_abi(abi, abi_file_path)
 
     def __getattr__(self, attribute):
@@ -71,9 +71,6 @@ def call_function(self, func: ABIFunction, args: tuple, kwargs: dict = {}):
             tx.contractInputsValues = None  # type: ignore
 
         for arg, input_type in zip(args, func.inputs):
-            if input_type.type == "address":
-                arg = self.tx_builder._resolve_address(arg)
-
             input_template = self.tx_builder.load_template(TemplateType.INPUT_TYPE)
             input_template.name = input_type.name
             input_template.type = input_type.type

diff --git a/bal_tools/safe_tx_builder/safe_tx_builder.py b/bal_tools/safe_tx_builder/safe_tx_builder.py
@@ -2,8 +2,6 @@
 from datetime import datetime, timezone
 import os
 
-from bal_addresses import AddrBook
-
 from .models import *
 from ..utils import is_address, chain_ids_by_name
 
@@ -46,10 +44,7 @@ def _initialize(
     ):
         self.chain_name = chain_name
         self.chain_id = str(chain_ids_by_name()[chain_name])
-        self.addr_book = AddrBook(
-            AddrBook.chain_names_by_id[int(self.chain_id)]
-        ).flatbook
-        self.safe_address = self._resolve_address(safe_address)
+        self.safe_address = safe_address
         self.version = version
         self.timestamp = timestamp if timestamp else datetime.now(timezone.utc)
         self.tx_builder_version = tx_builder_version
@@ -70,12 +65,6 @@ def load_template(
 
         return model.model_validate_json(file_content)
 
-    def _resolve_address(self, identifier: str) -> str:
-        if is_address(identifier):
-            return identifier
-
-        return self.addr_book[identifier]
-
     def _load_payload_metadata(self):
         self.base_payload.version = self.version
         self.base_payload.chainId = self.chain_id

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -2,7 +2,6 @@
 import json
 import os
 
-from bal_addresses import AddrBook
 from bal_tools.pools_gauges import BalPoolsGauges
 from bal_tools.subgraph import Subgraph
 from bal_tools.ecosystem import Aura
@@ -51,18 +50,8 @@ def aura(chain):
 
 
 @pytest.fixture(scope="module")
-def addr_book():
-    return AddrBook("mainnet").flatbook
-
-
-@pytest.fixture(scope="module")
-def msig_name():
-    return "multisigs/maxi_omni"
-
-
-@pytest.fixture(scope="module")
-def safe_tx_builder(msig_name) -> SafeTxBuilder:
-    return SafeTxBuilder(msig_name)
+def safe_tx_builder() -> SafeTxBuilder:
+    return SafeTxBuilder("0x10A19e7eE7d7F8a52822f6817de8ea18204F2e4f")
 
 
 @pytest.fixture(scope="module")

diff --git a/tests/test_safe_tx_builder.py b/tests/test_safe_tx_builder.py
@@ -7,29 +7,25 @@
 def test_safe_contract(
     safe_tx_builder: SafeTxBuilder,
     erc20_abi,
-    bribe_market_abi,
-    addr_book,
-    msig_name,
+    bribe_market_abi
 ):
     usdc_address = "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"
     bribe_market_address = "0x45Bc37b18E73A42A4a826357a8348cDC042cCBBc"
 
-    usdc = SafeContract("tokens/USDC", erc20_abi)
+    usdc = SafeContract(usdc_address, erc20_abi)
     usdc.approve(bribe_market_address, 100e18)
 
     bribe_market = SafeContract(bribe_market_address, bribe_market_abi)
     bribe_market.depositBribe(
         "0x0000000000000000000000000000000000000000000000000000000000000000",
-        "tokens/USDC",
+        usdc_address,
         1e18,
         0,
         2,
     )
 
     payload = safe_tx_builder.output_payload("tests/payload_outputs/bribe.json")
 
-    assert safe_tx_builder.safe_address == addr_book[msig_name]
-
     assert payload.transactions[0].to == usdc_address
     assert payload.transactions[0].contractMethod.name == "approve"
     assert payload.transactions[0].contractMethod.inputs[0].type == "address"
@@ -57,10 +53,12 @@ def test_safe_contract(
 
 
 def test_multiple_functions_with_same_name(bridge_abi):
-    builder = SafeTxBuilder("multisigs/dao")
+    usdc_address = "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"
+    dao_msig = "0x10A19e7eE7d7F8a52822f6817de8ea18204F2e4f"
+    builder = SafeTxBuilder(dao_msig)
     bridge = SafeContract("0x88ad09518695c6c3712AC10a214bE5109a655671", bridge_abi)
 
-    bridge.relayTokens("tokens/USDC", "multisigs/dao", int(1e18))
-    bridge.relayTokens("tokens/USDC", int(1e18))
-    
+    bridge.relayTokens(usdc_address, dao_msig, int(1e18))
+    bridge.relayTokens(usdc_address, int(1e18))
+
     builder.output_payload("tests/payload_outputs/multiple_names.json")