extend airdrop tests

proxy/plugin/solana_rest_api.py

@@ -8,7 +8,7 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 import json
 import unittest
 import rlp
@@ -25,7 +25,7 @@
 import base58
 import traceback
 import threading
-from .solana_rest_api_tools import EthereumAddress, create_account_with_seed, evm_loader_id, get_token_balance, \
+from .solana_rest_api_tools import EthereumAddress, create_account_with_seed, evm_loader_id, get_token_balance_or_airdrop, \
     getAccountInfo, solana_cli, call_signed, solana_url, call_emulated, \
     Trx,  EthereumError, create_collateral_pool_address, getTokenAddr, STORAGE_SIZE, neon_config_load
 from solana.rpc.commitment import Commitment, Confirmed
@@ -69,22 +69,8 @@ def __init__(self, client, signer, proxy_id):
 class EthereumModel:
     def __init__(self):
         # Initialize user account
-        res = solana_cli().call('config', 'get')
-        substr = "Keypair Path: "
-        path = ""
-        for line in res.splitlines():
-            if line.startswith(substr):
-                path = line[len(substr):].strip()
-        if path == "":
-            raise Exception("cannot get keypair path")
-
-        with open(path.strip(), mode='r') as file:
-            pk = (file.read())
-            nums = list(map(int, pk.strip("[] \n").split(',')))
-            nums = nums[0:32]
-            values = bytes(nums)
-            self.signer = sol_Account(values)
 
+        self.signer = self.get_solana_account()
         self.client = SolanaClient(solana_url)
 
         self.logs_db = LogDB(filename="local.db")
@@ -102,6 +88,26 @@ def __init__(self):
         neon_config_load(self)
         pass
 
+    @staticmethod
+    def get_solana_account() -> Optional[sol_Account]:
+        solana_account = None
+        res = solana_cli().call('config', 'get')
+        substr = "Keypair Path: "
+        path = ""
+        for line in res.splitlines():
+            if line.startswith(substr):
+                path = line[len(substr):].strip()
+        if path == "":
+            raise Exception("cannot get keypair path")
+
+        with open(path.strip(), mode='r') as file:
+            pk = (file.read())
+            nums = list(map(int, pk.strip("[] \n").split(',')))
+            nums = nums[0:32]
+            values = bytes(nums)
+            solana_account = sol_Account(values)
+        return solana_account
+
     def web3_clientVersion(self):
         neon_config_load(self)
         return self.neon_config_dict['web3_clientVersion']
@@ -150,7 +156,7 @@ def eth_getBalance(self, account, tag):
         """
         eth_acc = EthereumAddress(account)
         logger.debug('eth_getBalance: %s %s', account, eth_acc)
-        balance = get_token_balance(self.client, self.signer, evm_loader_id, eth_acc, self.signer.public_key())
+        balance = get_token_balance_or_airdrop(self.client, self.signer, evm_loader_id, eth_acc)
 
         return hex(balance*10**9)
 
@@ -478,12 +484,11 @@ def default(self, obj):
 
 
 class SolanaContractTests(unittest.TestCase):
+
     def setUp(self):
         self.model = EthereumModel()
         self.owner = '0xc1566af4699928fdf9be097ca3dc47ece39f8f8e'
         self.token1 = '0x49a449cd7fd8fbcf34d103d98f2c05245020e35b'
-#        self.assertEqual(self.getBalance(self.owner), 1000*10**18)
-#        self.assertEqual(self.getBalance(self.token1), 0)
 
     def getBalance(self, account):
         return int(self.model.eth_getBalance(account, 'latest'), 16)

proxy/plugin/solana_rest_api_tools.py

@@ -1238,32 +1238,33 @@ def create_and_mint_tkn_acc(client: SolanaClient, signer: SolanaAccount, eth_acc
         raise Exception("Create account error")
 
 
-def get_token_balance_impl(client: SolanaClient, account: str, eth_acc: EthereumAddress) -> [Optional[int], Optional[Union[Dict, str]]]:
-    token_account = get_associated_token_address(PublicKey(account), ETH_TOKEN_MINT_ID)
+def get_token_balance_gwei(client: SolanaClient, token_owner_acc: str, eth_acc: EthereumAddress) \
+                          -> [Optional[int], Optional[Union[Dict, str]]]:
+    token_account = get_associated_token_address(PublicKey(token_owner_acc), ETH_TOKEN_MINT_ID)
     rpc_response = client.get_token_account_balance(token_account, commitment=Confirmed)
     error = rpc_response.get('error')
     if error is None:
         balance = get_from_dict(rpc_response, "result", "value", "amount")
         if balance is None:
             return None, f"Failed to get token balance from: {rpc_response}, by eth account:" \
                          f" {eth_acc} aka: {token_account} at token: {ETH_TOKEN_MINT_ID}"
-        return balance, None
+        return int(balance), None
     return None, error
 
 
-def get_token_balance(client: SolanaClient, signer: SolanaAccount, evm_loader: str, eth_acc: EthereumAddress, base_account: PublicKey) -> int:
+def get_token_balance_or_airdrop(client: SolanaClient, signer: SolanaAccount, evm_loader: str, eth_acc: EthereumAddress) -> int:
 
-    account, nonce = ether2program(bytes(eth_acc).hex(), evm_loader, base_account)
+    account, nonce = ether2program(bytes(eth_acc).hex(), evm_loader, signer.public_key())
     logger.debug(f"Get balance for eth account: {eth_acc} aka: {account} at token: {ETH_TOKEN_MINT_ID}")
 
-    balance, error = get_token_balance_impl(client, account, eth_acc)
+    balance, error = get_token_balance_gwei(client, account, eth_acc)
     if error is None:
         return int(balance)
 
     if error.get("message") == SolanaErrors.AccountNotFound.value and NEW_USER_AIRDROP_AMOUNT > 0:
         logger.debug(f"Account not found:  {eth_acc} aka: {account} at token: {ETH_TOKEN_MINT_ID}")
         create_and_mint_tkn_acc(client, signer, eth_acc)
-        balance, error = get_token_balance_impl(client, account, eth_acc)
+        balance, error = get_token_balance_gwei(client, account, eth_acc)
         if error is None:
             return int(balance)
 

proxy/testing/test_airdropping_eth_accounts.py

@@ -1,53 +1,80 @@
 import unittest
+import os
+import solcx
+
 import eth_account
 import eth_typing
 import eth_utils
-import os
-from web3 import Web3
-import solcx
-import logging
+from eth_account.account import LocalAccount
+
+from web3 import Web3, eth as web3_eth
+from solana.rpc.api import Client as SolanaClient
 
-NEW_USER_AIRDROP_AMOUNT = int(os.environ.get("NEW_USER_AIRDROP_AMOUNT", "0"))
+from ..plugin.solana_rest_api import EthereumModel
+from ..plugin.solana_rest_api_tools import get_token_balance_gwei, EthereumAddress, ether2program
 
 
 class TestAirdroppingEthAccounts(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls) -> None:
+        cls._EVM_LOADER_ID = os.environ.get("EVM_LOADER")
+        new_user_airdrop_amount = int(os.environ.get("NEW_USER_AIRDROP_AMOUNT", "0"))
+        cls._EXPECTED_BALANCE_WEI = eth_utils.to_wei(new_user_airdrop_amount, 'ether')
+
         proxy_url = os.environ.get('PROXY_URL', 'http://localhost:9090/solana')
-        cls.proxy = Web3(Web3.HTTPProvider(proxy_url))
+        cls._web3 = Web3(Web3.HTTPProvider(proxy_url))
+        solana_url = os.environ.get("SOLANA_URL", "http://localhost:8899")
+        cls._solana_client = SolanaClient(solana_url)
+        cls._host_solana_account = EthereumModel.get_solana_account()
 
     def test_airdrop_on_get_balance(self):
-        account: eth_account.account.LocalAccount = eth_account.account.Account.create()
-        block_number: eth_typing.BlockNumber = self.proxy.eth.get_block_number()
-        actual_balance_wei = self.proxy.eth.get_balance(account.address, block_identifier=block_number)
-        expected_balance_wei = eth_utils.to_wei(NEW_USER_AIRDROP_AMOUNT, 'ether')
-        self.assertEqual(expected_balance_wei, actual_balance_wei)
+        account: LocalAccount = eth_account.account.Account.create()
+        block_number: eth_typing.BlockNumber = self._web3.eth.get_block_number()
+        actual_balance_wei = self._web3.eth.get_balance(account.address, block_identifier=block_number)
+        self.assertEqual(self._EXPECTED_BALANCE_WEI, actual_balance_wei)
 
     def test_airdrop_on_deploy(self):
-        contract_owner: eth_account.account.LocalAccount = self.proxy.eth.account.create()
-        compiled_sol = solcx.compile_source(self._CONTRACT_STORAGE_SOURCE)
+        contract_owner: LocalAccount = self._web3.eth.account.create()
+        contract = self._compile_and_deploy_contract(contract_owner, self._CONTRACT_STORAGE_SOURCE)
+        actual_balance_wei = self._get_balance_wei(contract.address)
+        self.assertEqual(self._EXPECTED_BALANCE_WEI, actual_balance_wei)
+
+    def test_airdrop_onto_wrapped_new_address(self):
+        contract_owner: LocalAccount = self._web3.eth.account.create()
+        contract = self._compile_and_deploy_contract(contract_owner, self._WRAPPER_CONTRACT_STORAGE_SOURCE)
+        nested_contract_address = contract.functions.getNested().call()
+        nested_actual_balance = self._get_balance_wei(nested_contract_address)
+        wrapper_actual_balance = self._get_balance_wei(contract.address)
+        self.assertEqual(self._EXPECTED_BALANCE_WEI, wrapper_actual_balance)
+        self.assertEqual(self._EXPECTED_BALANCE_WEI, nested_actual_balance)
+
+    def _compile_and_deploy_contract(self, contract_owner: LocalAccount, source: str) -> web3_eth.Contract:
+        compiled_sol = solcx.compile_source(source)
         contract_id, contract_interface = compiled_sol.popitem()
-        storage = self.proxy.eth.contract(abi=contract_interface['abi'], bytecode=contract_interface['bin'])
-        nonce = self.proxy.eth.get_transaction_count(contract_owner.address)
-        chain_id = self.proxy.eth.chain_id
-        trx_signed = self.proxy.eth.account.sign_transaction(
-            dict(nonce=nonce, chainId=chain_id, gas=987654321, gasPrice=0, to='', value=0, data=storage.bytecode),
+        contract = self._web3.eth.contract(abi=contract_interface['abi'], bytecode=contract_interface['bin'])
+        nonce = self._web3.eth.get_transaction_count(contract_owner.address)
+        chain_id = self._web3.eth.chain_id
+        trx_signed = self._web3.eth.account.sign_transaction(
+            dict(nonce=nonce, chainId=chain_id, gas=987654321, gasPrice=0, to='', value=0, data=contract.bytecode),
             contract_owner.key
         )
-        trx_hash = self.proxy.eth.send_raw_transaction(trx_signed.rawTransaction)
-        trx_receipt = self.proxy.eth.wait_for_transaction_receipt(trx_hash)
-        storage_contract = self.proxy.eth.contract(
+        trx_hash = self._web3.eth.send_raw_transaction(trx_signed.rawTransaction)
+        trx_receipt = self._web3.eth.wait_for_transaction_receipt(trx_hash)
+        contract = self._web3.eth.contract(
             address=trx_receipt.contractAddress,
-            abi=storage.abi
+            abi=contract.abi
         )
-        actual_balance_wei = self.proxy.eth.get_balance(storage_contract.address, block_identifier="latest")
-        expected_balance_wei = eth_utils.to_wei(NEW_USER_AIRDROP_AMOUNT, 'ether')
+        return contract
 
-        owner_balance = self.proxy.eth.get_balance(contract_owner.address, block_identifier="latest")
-
-        self.assertEqual(expected_balance_wei, owner_balance)
-        self.assertEqual(expected_balance_wei, actual_balance_wei)
+    def _get_balance_wei(self, eth_acc: str) -> int:
+        pub_key = self._host_solana_account.public_key()
+        token_owner_account, nonce = ether2program(eth_acc, self._EVM_LOADER_ID, pub_key)
+        balance, error = get_token_balance_gwei(self._solana_client, token_owner_account, EthereumAddress(eth_acc))
+        self.assertIsNone(error)
+        self.assertIsNotNone(balance)
+        self.assertIsInstance(balance, int)
+        return balance * eth_utils.denoms.gwei
 
     _CONTRACT_STORAGE_SOURCE = '''pragma solidity >=0.7.0 <0.9.0;
         contract Storage {
@@ -61,3 +88,18 @@ def test_airdrop_on_deploy(self):
         }
     '''
 
+    _WRAPPER_CONTRACT_STORAGE_SOURCE = '''
+        // SPDX-License-Identifier: GPL-3.0
+        pragma solidity >=0.7.0 <0.9.0;
+        contract Wrapper {
+            address private nested_address;
+            constructor() {
+                Nested nested = new Nested();
+                nested_address = address(nested);
+            }
+            function getNested() public view returns (address) {
+                return nested_address;
+            }
+        }
+        contract Nested {}
+    '''