Skip to content

Commit

Permalink
feat(fw): also test system contract deployments with balance
Browse files Browse the repository at this point in the history
  • Loading branch information
danceratopz committed Jan 27, 2025
1 parent 51a889a commit 5fd12b5
Showing 1 changed file with 31 additions and 5 deletions.
36 changes: 31 additions & 5 deletions src/ethereum_test_tools/utility/generators.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,13 @@ class DeploymentTestType(Enum):
DEPLOY_AFTER_FORK = "deploy_after_fork"


class ContractAddressHasBalance(Enum):
"""Represents whether the target deployment test has a balance before deployment."""

ZERO_BALANCE = "zero_balance"
NONZERO_BALANCE = "nonzero_balance"


class SystemContractDeployTestFunction(Protocol):
"""
Represents a function to be decorated with the `generate_system_contract_deploy_test`
Expand Down Expand Up @@ -61,9 +68,17 @@ def generate_system_contract_deploy_test(
"""
Generate a test that verifies the correct deployment of a system contract.
Generates two tests:
- One that deploys the contract before the fork.
- One that deploys the contract after the fork.
Generates four test cases:
| before/after fork | has balance |
------------------------------------|-------------------|-------------|
`deploy_before_fork-nonzero_balance`| before | True |
`deploy_before_fork-zero_balance` | before | False |
`deploy_after_fork-nonzero_balance` | after | True |
`deploy_after_fork-zero_balance` | after | False |
where `has balance` refers to whether the contract address has a non-zero balance before
deployment, or not.
Args:
fork (Fork): The fork to test.
Expand All @@ -89,6 +104,14 @@ def generate_system_contract_deploy_test(
deployer_address = deploy_tx.sender

def decorator(func: SystemContractDeployTestFunction):
@pytest.mark.parametrize(
"has_balance",
[
pytest.param(ContractAddressHasBalance.NONZERO_BALANCE),
pytest.param(ContractAddressHasBalance.ZERO_BALANCE),
],
ids=lambda x: x.name.lower(),
)
@pytest.mark.parametrize(
"test_type",
[
Expand All @@ -101,6 +124,7 @@ def decorator(func: SystemContractDeployTestFunction):
@pytest.mark.valid_at_transition_to(fork.name())
def wrapper(
blockchain_test: BlockchainTestFiller,
has_balance: ContractAddressHasBalance,
pre: Alloc,
test_type: DeploymentTestType,
fork: Fork,
Expand Down Expand Up @@ -131,11 +155,11 @@ def wrapper(
timestamp=15_001,
),
]

balance = 1 if has_balance == ContractAddressHasBalance.NONZERO_BALANCE else 0
pre[expected_deploy_address] = Account(
code=b"", # Remove the code that is automatically allocated on the fork
nonce=0,
balance=0,
balance=balance,
)
pre[deployer_address] = Account(
balance=deployer_required_balance,
Expand All @@ -151,12 +175,14 @@ def wrapper(
post[expected_deploy_address] = Account(
code=expected_code,
nonce=1,
balance=balance,
)
else:
post[expected_deploy_address] = Account(
storage=expected_system_contract_storage,
code=expected_code,
nonce=1,
balance=balance,
)
post[deployer_address] = Account(
nonce=1,
Expand Down

0 comments on commit 5fd12b5

Please sign in to comment.