Fix: Fee token and Account improvements

starknet_devnet/account.py

@@ -9,7 +9,6 @@
 from starkware.starknet.services.api.contract_class import ContractClass
 from starkware.starknet.core.os.contract_address.contract_address import calculate_contract_address_from_hash
 from starkware.starknet.storage.starknet_storage import StorageLeaf
-from starkware.starknet.testing.starknet import Starknet
 from starkware.starknet.testing.contract import StarknetContract
 from starkware.python.utils import to_bytes
 
@@ -26,7 +25,8 @@ class Account:
     HASH = 3234970678029762354735267567433689214900679403476863445247436772798892968339
     HASH_BYTES = to_bytes(HASH)
 
-    def __init__(self, private_key: int, public_key: int, initial_balance: int):
+    def __init__(self, starknet_wrapper, private_key: int, public_key: int, initial_balance: int):
+        self.starknet_wrapper = starknet_wrapper
         self.private_key = private_key
         self.public_key = public_key
 
@@ -56,13 +56,15 @@ def to_json(self):
             "address": hex(self.address)
         }
 
-    async def deploy(self, starknet: Starknet) -> StarknetContract:
+    async def deploy(self) -> StarknetContract:
         """Deploy this account."""
+        starknet = self.starknet_wrapper.starknet
+        contract_class = Account.get_contract_class()
         account_carried_state = starknet.state.state.contract_states[self.address]
         account_state = account_carried_state.state
         assert not account_state.initialized
 
-        starknet.state.state.contract_definitions[Account.HASH_BYTES] = Account.get_contract_class()
+        starknet.state.state.contract_definitions[Account.HASH_BYTES] = contract_class
 
         newly_deployed_account_state = await ContractState.create(
             contract_hash=Account.HASH_BYTES,
@@ -85,9 +87,11 @@ async def deploy(self, starknet: Starknet) -> StarknetContract:
         fee_token_storage_updates[balance_address] = StorageLeaf(initial_balance_uint256.low)
         fee_token_storage_updates[balance_address + 1] = StorageLeaf(initial_balance_uint256.high)
 
-        return StarknetContract(
+        contract =  StarknetContract(
             state=starknet.state,
-            abi=Account.get_contract_class().abi,
+            abi=contract_class.abi,
             contract_address=self.address,
             deploy_execution_info=None
         )
+
+        self.starknet_wrapper.store_contract(self.address, contract, contract_class)

starknet_devnet/accounts.py

@@ -0,0 +1,70 @@
+"""
+Class representing list of predefined accounts
+"""
+
+import random
+import sys
+
+from typing import List
+from starkware.crypto.signature.signature import private_to_stark_key
+from .account import Account
+
+class Accounts:
+    """Accounts wrapper"""
+    list: List[Account] = []
+
+    def __init__(self, starknet_wrapper):
+        self.starknet_wrapper = starknet_wrapper
+        self.__initial_balance = None
+        self.__seed = None
+
+    def __getitem__(self, index):
+        return self.list[index]
+
+    async def deploy(self):
+        """deploy listed accounts"""
+        for account in self.list:
+            await account.deploy()
+
+    def add(self, account):
+        """append account to list"""
+        self.list.append(account)
+        return account
+
+    def generate(self, n_accounts: int, initial_balance: int, seed: int):
+        """Generates accounts without deploying them"""
+        random_generator = random.Random()
+        self.__initial_balance = initial_balance
+        self.__seed = seed
+
+        if seed is None:
+            seed = random_generator.getrandbits(32)
+        random_generator.seed(seed)
+
+        for _ in range(n_accounts):
+            private_key = random_generator.getrandbits(128)
+            public_key = private_to_stark_key(private_key)
+
+            self.add(Account(
+                self.starknet_wrapper,
+                private_key=private_key,
+                public_key=public_key,
+                initial_balance=initial_balance
+            ))
+
+    def print(self):
+        """stdout accounts list"""
+        for idx, account in enumerate(self):
+            print(f"Account #{idx}")
+            print(f"Address: {hex(account.address)}")
+            print(f"Public key: {hex(account.public_key)}")
+            print(f"Private key: {hex(account.private_key)}\n")
+
+        print(f"Initial balance of each account: {self.__initial_balance} WEI")
+        print("Seed to replicate this account sequence:", self.__seed)
+        print(
+            "WARNING: Use these accounts and their keys ONLY for local testing. "
+            "DO NOT use them on mainnet or other live networks because you will LOSE FUNDS.\n",
+            file=sys.stderr
+        )
+        sys.stdout.flush()

starknet_devnet/fee_token.py

@@ -11,7 +11,6 @@
 from starkware.python.utils import to_bytes
 from starkware.starknet.compiler.compile import get_selector_from_name
 from starknet_devnet.util import Uint256
-from .contract_wrapper import ContractWrapper
 
 class FeeToken:
     """Wrapper of token for charging fees."""
@@ -32,8 +31,8 @@ class FeeToken:
 
     contract: StarknetContract = None
 
-    def __init__(self, startknet_wrapper):
-        self.startknet_wrapper = startknet_wrapper
+    def __init__(self, starknet_wrapper):
+        self.starknet_wrapper = starknet_wrapper
 
     @classmethod
     def get_contract_class(cls):
@@ -44,13 +43,14 @@ def get_contract_class(cls):
 
     async def deploy(self):
         """Deploy token contract for charging fees."""
-        starknet = self.startknet_wrapper.starknet
+        starknet = self.starknet_wrapper.starknet
+        contract_class = FeeToken.get_contract_class()
 
         fee_token_carried_state = starknet.state.state.contract_states[FeeToken.ADDRESS]
         fee_token_state = fee_token_carried_state.state
         assert not fee_token_state.initialized
 
-        starknet.state.state.contract_definitions[FeeToken.HASH_BYTES] = FeeToken.get_contract_class()
+        starknet.state.state.contract_definitions[FeeToken.HASH_BYTES] = contract_class
         newly_deployed_fee_token_state = await ContractState.create(
             contract_hash=FeeToken.HASH_BYTES,
             storage_commitment_tree=fee_token_state.storage_commitment_tree
@@ -73,7 +73,7 @@ async def deploy(self):
             deploy_execution_info=None
         )
 
-        self.startknet_wrapper.contracts.store(FeeToken.ADDRESS, ContractWrapper(self.contract, self.get_contract_class()))
+        self.starknet_wrapper.store_contract(FeeToken.ADDRESS, self.contract, contract_class)
 
     async def get_balance(self, address: int) -> int:
         """Return the balance of the contract under `address`."""
@@ -115,7 +115,7 @@ async def mint(self, to_address: int, amount: int, lite: bool):
             ).invoke()
         else:
             transaction = self.get_mint_transaction(to_address, amount_uint256)
-            _, tx_hash_int, _ = await self.startknet_wrapper.invoke(transaction)
+            _, tx_hash_int, _ = await self.starknet_wrapper.invoke(transaction)
             tx_hash = hex(tx_hash_int)
 
         return tx_hash

starknet_devnet/server.py

@@ -36,11 +36,12 @@ async def initialize_starknet():
 def generate_accounts(args):
     """Generate accounts """
     if args.accounts:
-        state.generate_accounts(
+        state.starknet_wrapper.accounts.generate(
             n_accounts=args.accounts,
             initial_balance=args.initial_balance,
             seed=args.seed
         )
+        state.starknet_wrapper.accounts.print()
 
 def set_dump_options(args):
     """Assign dumping options from args to state."""

starknet_devnet/starknet_wrapper.py

@@ -24,7 +24,7 @@
 from starkware.starknet.testing.contract import StarknetContract
 from starkware.starknet.testing.objects import FunctionInvocation
 
-from .account import Account
+from .accounts import Accounts
 from .fee_token import FeeToken
 from .general_config import DEFAULT_GENERAL_CONFIG
 from .origin import NullOrigin, Origin
@@ -66,14 +66,11 @@ def __init__(self, config: DevnetConfig):
         self.contracts = DevnetContracts(self.origin)
         self.l1l2 = DevnetL1L2()
         self.transactions = DevnetTransactions(self.origin)
-        self.__starknet = None
-        self.starknet = None # Temporary hotfix will be cleaned with async fix PR
+        self.starknet: Starknet = None
         self.__current_carried_state = None
         self.__initialized = False
         self.fee_token = FeeToken(self)
-
-        self.accounts: List[Account] = []
-        """List of predefined accounts"""
+        self.accounts = Accounts(self)
 
     @staticmethod
     def load(path: str) -> "StarknetWrapper":
@@ -87,7 +84,7 @@ async def initialize(self):
             starknet = await self.__get_starknet()
 
             await self.fee_token.deploy()
-            await self.__deploy_accounts()
+            await self.accounts.deploy()
 
             await self.__preserve_current_state(starknet.state.state)
             self.__initialized = True
@@ -100,10 +97,10 @@ async def __get_starknet(self):
         """
         Returns the underlying Starknet instance, creating it first if necessary.
         """
-        if not self.__starknet:
-            self.__starknet = await Starknet.empty(general_config=DEFAULT_GENERAL_CONFIG)
-        self.starknet = self.__starknet
-        return self.__starknet
+        if not self.starknet:
+            self.starknet = await Starknet.empty(general_config=DEFAULT_GENERAL_CONFIG)
+
+        return self.starknet
 
     async def get_state(self):
         """
@@ -145,6 +142,10 @@ async def __get_state_root(self):
         state = await self.get_state()
         return state.state.shared_state.contract_states.root
 
+    def store_contract(self, address, contract, contract_class):
+        """Store the provided data sa wrapped contract"""
+        self.contracts.store(address, ContractWrapper(contract, contract_class))
+
     async def __store_transaction(
         self, transaction: DevnetTransaction, tx_hash: int,
         state_update: Dict, error_message: str=None
@@ -171,12 +172,6 @@ async def __store_transaction(
 
         self.transactions.store(tx_hash, transaction)
 
-    async def __deploy_accounts(self):
-        starknet = await self.__get_starknet()
-        for account in self.accounts:
-            contract = await account.deploy(starknet)
-            self.contracts.store(account.address, ContractWrapper(contract, Account.get_contract_class()))
-
     def set_config(self, config: DevnetConfig):
         """
         Sets the configuration of the devnet.
@@ -330,8 +325,6 @@ async def call(self, transaction: InvokeFunction):
 
         return { "result": adapted_result }
 
-
-
     async def __register_new_contracts(self, internal_calls: List[Union[FunctionInvocation, CallInfo]], tx_hash: int):
         for internal_call in internal_calls:
             if internal_call.entry_point_type == EntryPointType.CONSTRUCTOR:

starknet_devnet/state.py

@@ -2,12 +2,6 @@
 Global state singletone
 """
 
-import random
-import sys
-
-from starkware.crypto.signature.signature import private_to_stark_key
-
-from .account import Account
 from .dump import Dumper
 from .starknet_wrapper import StarknetWrapper, DevnetConfig
 
@@ -34,42 +28,4 @@ def load(self, load_path: str):
         """Loads starknet wrapper from path"""
         self.__set_starknet_wrapper(StarknetWrapper.load(load_path))
 
-    def generate_accounts(self, n_accounts: int, initial_balance: int, seed: int):
-        """Generates accounts without deploying them"""
-        random_generator = random.Random()
-
-        if seed is None:
-            seed = random_generator.getrandbits(32)
-        random_generator.seed(seed)
-
-        accounts = []
-        for i in range(n_accounts):
-            private_key = random_generator.getrandbits(128)
-            public_key = private_to_stark_key(private_key)
-
-            account = Account(
-                private_key=private_key,
-                public_key=public_key,
-                initial_balance=initial_balance
-            )
-            accounts.append(account)
-
-            print(f"Account #{i}")
-            print(f"Address: {hex(account.address)}")
-            print(f"Public key: {hex(account.public_key)}")
-            print(f"Private key: {hex(account.private_key)}")
-            print()
-
-        self.starknet_wrapper.accounts = accounts
-
-        print(f"Initial balance of each account: {initial_balance} WEI")
-        print("Seed to replicate this account sequence:", seed)
-        print(
-            "WARNING: Use these accounts and their keys ONLY for local testing. "
-            "DO NOT use them on mainnet or other live networks because you will LOSE FUNDS.",
-            file=sys.stderr
-        )
-        print()
-        sys.stdout.flush()
-
 state = State()