Refactor predeployed contract wrapper classes

starknet_devnet/account.py

@@ -6,14 +6,14 @@
     calculate_contract_address_from_hash,
 )
 from starkware.starknet.public.abi import get_selector_from_name
-from starkware.starknet.testing.contract import StarknetContract
 from starkware.starknet.testing.starknet import Starknet
 
 from starknet_devnet.account_util import set_balance
 from starknet_devnet.contract_class_wrapper import ContractClassWrapper
+from starknet_devnet.predeployed_contract_wrapper import PredeployedContractWrapper
 
 
-class Account:
+class Account(PredeployedContractWrapper):
     """Account contract wrapper."""
 
     # pylint: disable=too-many-arguments
@@ -50,14 +50,8 @@ def to_json(self):
             "address": hex(self.address),
         }
 
-    async def deploy(self) -> StarknetContract:
-        """Deploy this account and set its balance."""
+    async def _mimic_constructor(self):
         starknet: Starknet = self.starknet_wrapper.starknet
-        contract_class = self.contract_class
-        await starknet.state.state.set_contract_class(
-            self.class_hash_bytes, contract_class
-        )
-        await starknet.state.state.deploy_contract(self.address, self.class_hash_bytes)
 
         await starknet.state.state.set_storage_at(
             self.address, get_selector_from_name("Account_public_key"), self.public_key

starknet_devnet/fee_token.py

@@ -7,17 +7,17 @@
 from starkware.starknet.compiler.compile import get_selector_from_name
 from starkware.starknet.services.api.contract_class import ContractClass
 from starkware.starknet.services.api.gateway.transaction import InvokeFunction
-from starkware.starknet.testing.contract import StarknetContract
 from starkware.starknet.testing.starknet import Starknet
 
 from starknet_devnet.account_util import get_execute_args
 from starknet_devnet.chargeable_account import ChargeableAccount
 from starknet_devnet.constants import SUPPORTED_TX_VERSION
+from starknet_devnet.predeployed_contract_wrapper import PredeployedContractWrapper
 from starknet_devnet.sequencer_api_utils import InternalInvokeFunction
 from starknet_devnet.util import Uint256, str_to_felt
 
 
-class FeeToken:
+class FeeToken(PredeployedContractWrapper):
     """Wrapper of token for charging fees."""
 
     CONTRACT_CLASS: ContractClass = None  # loaded lazily
@@ -33,10 +33,10 @@ class FeeToken:
     SYMBOL = "ETH"
     NAME = "ether"
 
-    contract: StarknetContract = None
-
     def __init__(self, starknet_wrapper):
         self.starknet_wrapper = starknet_wrapper
+        self.address = self.ADDRESS
+        self.class_hash_bytes = self.HASH_BYTES
 
     @classmethod
     def get_contract_class(cls):
@@ -47,22 +47,13 @@ def get_contract_class(cls):
             )
         return cls.CONTRACT_CLASS
 
-    async def deploy(self):
-        """Deploy token contract for charging fees."""
-        starknet: Starknet = self.starknet_wrapper.starknet
-        contract_class = FeeToken.get_contract_class()
-
-        await starknet.state.state.set_contract_class(
-            FeeToken.HASH_BYTES, contract_class
-        )
-
-        # pylint: disable=protected-access
-        starknet.state.state.cache._class_hash_writes[
-            FeeToken.ADDRESS
-        ] = FeeToken.HASH_BYTES
-        # replace with await starknet.state.state.deploy_contract
+    @property
+    def contract_class(self) -> ContractClass:
+        """Same as `get_contract_class`, used by `PredeployedContractWrapper` parent"""
+        return self.get_contract_class()
 
-        # mimic constructor
+    async def _mimic_constructor(self):
+        starknet: Starknet = self.starknet_wrapper.starknet
         await starknet.state.state.set_storage_at(
             FeeToken.ADDRESS,
             get_selector_from_name("ERC20_name"),
@@ -84,13 +75,6 @@ async def deploy(self):
             ChargeableAccount.ADDRESS,
         )
 
-        self.contract = StarknetContract(
-            state=starknet.state,
-            abi=contract_class.abi,
-            contract_address=FeeToken.ADDRESS,
-            deploy_call_info=None,
-        )
-
     async def get_balance(self, address: int) -> int:
         """Return the balance of the contract under `address`."""
         response = await self.contract.balanceOf(address).call()

starknet_devnet/predeployed_contract_wrapper.py

@@ -0,0 +1,55 @@
+"""Parent for predeployed contract wrapper classes"""
+from abc import ABC
+
+from starkware.starknet.services.api.contract_class import ContractClass
+from starkware.starknet.testing.contract import StarknetContract
+from starkware.starknet.testing.starknet import Starknet
+
+
+class PredeployedContractWrapper(ABC):
+    """Parent class for all predeployed contract wrapper classes"""
+
+    # Cannot import it because of circular imports
+    # from .starknet_wrapper import StarknetWrapper
+
+    starknet_wrapper: "StarknetWrapper"
+    address: int
+    class_hash_bytes: bytes
+    contract_class: ContractClass
+
+    # Value will be set by deploy
+    contract: StarknetContract
+
+    async def _mimic_constructor(self):
+        raise NotImplementedError()
+
+    async def deploy(self):
+        """Deploy the contract wrapper to devnet"""
+        starknet: Starknet = self.starknet_wrapper.starknet
+
+        await starknet.state.state.set_contract_class(
+            self.class_hash_bytes, self.contract_class
+        )
+
+        # pylint: disable=protected-access
+        starknet: Starknet = self.starknet_wrapper.starknet
+        starknet.state.state.cache._class_hash_writes[
+            self.address
+        ] = self.class_hash_bytes
+        # replace with await starknet.state.state.deploy_contract
+        # await starknet.state.state.deploy_contract(self.address, self.class_hash_bytes)
+        # For now, it fails for fee token since the address is the same as the
+        # ETH Token, see:
+        # https://starkscan.co/token/0x049d36570d4e46f48e99674bd3fcc84644ddd6b96f7c741b1562b82f9e004dc7:
+        # Requested contract address
+        # 0x49d36570d4e46f48e99674bd3fcc84644ddd6b96f7c741b1562b82f9e004dc7
+        # is unavailable for deployment
+
+        await self._mimic_constructor()
+
+        self.contract = StarknetContract(
+            state=starknet.state,
+            abi=self.contract_class.abi,
+            contract_address=self.address,
+            deploy_call_info=None,
+        )

starknet_devnet/udc.py

@@ -3,10 +3,11 @@
 from starkware.python.utils import to_bytes
 from starkware.solidity.utils import load_nearby_contract
 from starkware.starknet.services.api.contract_class import ContractClass
-from starkware.starknet.testing.starknet import Starknet
 
+from starknet_devnet.predeployed_contract_wrapper import PredeployedContractWrapper
 
-class UDC:
+
+class UDC(PredeployedContractWrapper):
     """Universal deployer contract wrapper class"""
 
     CONTRACT_CLASS: ContractClass = None  # loaded lazily
@@ -23,6 +24,8 @@ class UDC:
 
     def __init__(self, starknet_wrapper):
         self.starknet_wrapper = starknet_wrapper
+        self.address = self.ADDRESS
+        self.class_hash_bytes = self.HASH_BYTES
 
     @classmethod
     def get_contract_class(cls):
@@ -33,13 +36,10 @@ def get_contract_class(cls):
             )
         return cls.CONTRACT_CLASS
 
-    async def deploy(self):
-        """Deploy token contract for charging fees."""
-        starknet: Starknet = self.starknet_wrapper.starknet
-        contract_class = UDC.get_contract_class()
-
-        await starknet.state.state.set_contract_class(UDC.HASH_BYTES, contract_class)
+    @property
+    def contract_class(self) -> ContractClass:
+        """Same as `get_contract_class`, used by `PredeployedContractWrapper` parent"""
+        return self.get_contract_class()
 
-        # pylint: disable=protected-access
-        starknet.state.state.cache._class_hash_writes[UDC.ADDRESS] = UDC.HASH_BYTES
-        # replace with await starknet.state.state.deploy_contract
+    async def _mimic_constructor(self):
+        pass