Skip to content

Commit 5dc2e5d

Browse files
authored
feat(execute): Add --skip-cleanup flag to execute (#2114)
1 parent ff6aca4 commit 5dc2e5d

File tree

1 file changed

+40
-25
lines changed

1 file changed

+40
-25
lines changed

src/pytest_plugins/execute/pre_alloc.py

Lines changed: 40 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,13 @@ def pytest_addoption(parser):
119119
type=int,
120120
help="The default amount of wei to fund each EOA in each test with.",
121121
)
122+
pre_alloc_group.addoption(
123+
"--skip-cleanup",
124+
action="store_true",
125+
dest="skip_cleanup",
126+
default=False,
127+
help="Skip cleanup phase after each test.",
128+
)
122129

123130

124131
@pytest.hookimpl(trylast=True)
@@ -143,6 +150,12 @@ def address_stubs(request) -> AddressStubs | None:
143150
return request.config.getoption("address_stubs", None)
144151

145152

153+
@pytest.fixture(scope="session")
154+
def skip_cleanup(request) -> bool:
155+
"""Return whether to skip cleanup phase after each test."""
156+
return request.config.getoption("skip_cleanup")
157+
158+
146159
@pytest.fixture(scope="session")
147160
def eoa_iterator(request) -> Iterator[EOA]:
148161
"""Return an iterator that generates EOAs."""
@@ -522,6 +535,7 @@ def pre(
522535
eoa_fund_amount_default: int,
523536
default_gas_price: int,
524537
address_stubs: AddressStubs | None,
538+
skip_cleanup: bool,
525539
request: pytest.FixtureRequest,
526540
) -> Generator[Alloc, None, None]:
527541
"""Return default pre allocation for all tests (Empty alloc)."""
@@ -544,31 +558,32 @@ def pre(
544558
# Yield the pre-alloc for usage during the test
545559
yield pre
546560

547-
# Refund all EOAs (regardless of whether the test passed or failed)
548-
refund_txs = []
549-
for idx, eoa in enumerate(pre._funded_eoa):
550-
remaining_balance = eth_rpc.get_balance(eoa)
551-
eoa.nonce = Number(eth_rpc.get_transaction_count(eoa))
552-
refund_gas_limit = 21_000
553-
tx_cost = refund_gas_limit * default_gas_price
554-
if remaining_balance < tx_cost:
555-
continue
556-
refund_tx = Transaction(
557-
sender=eoa,
558-
to=sender_key,
559-
gas_limit=21_000,
560-
gas_price=default_gas_price,
561-
value=remaining_balance - tx_cost,
562-
).with_signature_and_sender()
563-
refund_tx.metadata = TransactionTestMetadata(
564-
test_id=request.node.nodeid,
565-
phase="cleanup",
566-
action="refund_from_eoa",
567-
target=eoa.label,
568-
tx_index=idx,
569-
)
570-
refund_txs.append(refund_tx)
571-
eth_rpc.send_wait_transactions(refund_txs)
561+
if not skip_cleanup:
562+
# Refund all EOAs (regardless of whether the test passed or failed)
563+
refund_txs = []
564+
for idx, eoa in enumerate(pre._funded_eoa):
565+
remaining_balance = eth_rpc.get_balance(eoa)
566+
eoa.nonce = Number(eth_rpc.get_transaction_count(eoa))
567+
refund_gas_limit = 21_000
568+
tx_cost = refund_gas_limit * default_gas_price
569+
if remaining_balance < tx_cost:
570+
continue
571+
refund_tx = Transaction(
572+
sender=eoa,
573+
to=sender_key,
574+
gas_limit=21_000,
575+
gas_price=default_gas_price,
576+
value=remaining_balance - tx_cost,
577+
).with_signature_and_sender()
578+
refund_tx.metadata = TransactionTestMetadata(
579+
test_id=request.node.nodeid,
580+
phase="cleanup",
581+
action="refund_from_eoa",
582+
target=eoa.label,
583+
tx_index=idx,
584+
)
585+
refund_txs.append(refund_tx)
586+
eth_rpc.send_wait_transactions(refund_txs)
572587

573588
# Record the ending balance of the sender
574589
sender_test_ending_balance = eth_rpc.get_balance(sender_key)

0 commit comments

Comments
 (0)