feat: add --gas flag to ape test command to output gas reports after tests

.gitignore

@@ -115,6 +115,8 @@ venv.bak/
 .dmypy.json
 dmypy.json
 
+**/.DS_Store
+
 # setuptools-scm
 version.py
 

docs/userguides/testing.md

@@ -266,3 +266,37 @@ When you exit a provider's context, Ape **does not** disconnect the provider.
 When you re-enter that provider's context, Ape uses the previously-connected provider.
 At the end of the tests, Ape disconnects all the providers.
 Thus, you can enter and exit a provider's context as much as you need in tests.
+
+## Gas Reporting
+
+To include a gas report at the end of your tests, you can use the `--gas` flag.
+**NOTE**: This feature requires using a provider with tracing support, such as [ape-hardhat](https://github.com/ApeWorX/ape-hardhat).
+
+```bash
+ape test --network ethereum:local:hardhat --gas
+```
+
+At the end of test suite, you will see tables such as:
+
+```sh
+                            FundMe.vy Gas
+
+  Method           Times called    Min.    Max.    Mean   Median
+ ────────────────────────────────────────────────────────────────
+  fund                        8   57198   91398   82848    91398
+  withdraw                    2   28307   38679   33493    33493
+  changeOnStatus              2   23827   45739   34783    34783
+  getSecret                   1   24564   24564   24564    24564
+
+                  Transferring ETH Gas
+
+  Method     Times called   Min.   Max.   Mean   Median
+ ───────────────────────────────────────────────────────
+  to:test0              2   2400   9100   5750     5750
+
+                     TestContract.vy Gas
+
+  Method      Times called    Min.    Max.    Mean   Median
+ ───────────────────────────────────────────────────────────
+  setNumber              1   51021   51021   51021    51021
+```

docs/userguides/transactions.md

@@ -246,17 +246,15 @@ receipt.show_gas_report()
 It will output tables of contracts and methods with gas usages that look like this:
 
 ```bash
-                           DAI Gas
-┏━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━┳━━━━━━━┳━━━━━━━━┓
-┃ Method    ┃ Times called ┃ Min.  ┃ Max.  ┃ Mean  ┃ Median ┃
-┡━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━╇━━━━━━━╇━━━━━━━━┩
-│ balanceOf │ 4            │ 1302  │ 1302  │ 1302  │ 1302   │
-│ transfer  │ 5            │ 6974  │ 30374 │ 20174 │ 26174  │
-│ allowance │ 1            │ 1377  │ 1377  │ 1377  │ 1377   │
-│ approve   │ 1            │ 22414 │ 22414 │ 22414 │ 22414  │
-│ burn      │ 1            │ 11946 │ 11946 │ 11946 │ 11946  │
-│ mint      │ 1            │ 25845 │ 25845 │ 25845 │ 25845  │
-└───────────┴──────────────┴───────┴───────┴───────┴────────┘
+                            DAI Gas
+
+  Method           Times called    Min.    Max.    Mean   Median
+ ────────────────────────────────────────────────────────────────
+  balanceOf                   4   1302    13028   1302    1302
+  allowance                   2   1377    1377    1337    1337
+│ approve                     1   22414   22414   22414   22414
+│ burn                        1   11946   11946   11946   11946
+│ mint                        1   25845   25845   25845   25845
 ```
 
 ## Estimate Gas Cost

src/ape/api/transactions.py

@@ -174,6 +174,10 @@ def confirm_transaction(cls, value):
 
         return value
 
+    @property
+    def call_tree(self) -> Optional[Any]:
+        return None
+
     @property
     def failed(self) -> bool:
         """

src/ape/pytest/fixtures.py

@@ -1,23 +1,38 @@
-from typing import Iterator, List
+from typing import Dict, Iterator, List, Optional
 
 import pytest
+from _pytest.config import Config as PytestConfig
+from evm_trace.gas import merge_reports
 
-from ape.api import TestAccountAPI
-from ape.exceptions import ProviderNotConnectedError
+from ape.api import ReceiptAPI, TestAccountAPI
 from ape.logging import logger
 from ape.managers.chain import ChainManager
 from ape.managers.networks import NetworkManager
 from ape.managers.project import ProjectManager
-from ape.utils import ManagerAccessMixin
+from ape.types import GasReport, SnapshotID
+from ape.utils import CallTraceParser, ManagerAccessMixin, allow_disconnected, cached_property
 
 
 class PytestApeFixtures(ManagerAccessMixin):
     # NOTE: Avoid including links, markdown, or rst in method-docs
     # for fixtures, as they are used in output from the command
     # `ape test -q --fixture` (`pytest -q --fixture`).
 
-    def __init__(self):
-        self._warned_for_unimplemented_snapshot = False
+    _warned_for_unimplemented_snapshot = False
+    pytest_config: PytestConfig
+    receipt_capture: "ReceiptCapture"
+
+    def __init__(self, pytest_config, receipt_capture: "ReceiptCapture"):
+        self.pytest_config = pytest_config
+        self.receipt_capture = receipt_capture
+
+    @cached_property
+    def _using_traces(self) -> bool:
+        return (
+            self.network_manager.provider is not None
+            and self.provider.is_connected
+            and self.provider.supports_tracing
+        )
 
     @pytest.fixture(scope="session")
     def accounts(self) -> List[TestAccountAPI]:
@@ -54,32 +69,114 @@ def project(self) -> ProjectManager:
     def _isolation(self) -> Iterator[None]:
         """
         Isolation logic used to implement isolation fixtures for each pytest scope.
+        When tracing support is available, will also assist in capturing receipts.
         """
-        snapshot_id = None
-        try:
-            snapshot_id = self.chain_manager.snapshot()
-        except ProviderNotConnectedError:
-            logger.warning("Provider became disconnected mid-test.")
 
-        except NotImplementedError:
-            if not self._warned_for_unimplemented_snapshot:
-                logger.warning(
-                    "The connected provider does not support snapshotting. "
-                    "Tests will not be completely isolated."
-                )
-                self._warned_for_unimplemented_snapshot = True
+        snapshot_id = self._snapshot()
 
-        yield
+        if self._using_traces:
+            with self.receipt_capture:
+                yield
+        else:
+            yield
 
-        if snapshot_id is not None and snapshot_id in self.chain_manager._snapshots:
-            try:
-                self.chain_manager.restore(snapshot_id)
-            except ProviderNotConnectedError:
-                logger.warning("Provider became disconnected mid-test.")
+        if snapshot_id:
+            self._restore(snapshot_id)
 
     # isolation fixtures
     _session_isolation = pytest.fixture(_isolation, scope="session")
     _package_isolation = pytest.fixture(_isolation, scope="package")
     _module_isolation = pytest.fixture(_isolation, scope="module")
     _class_isolation = pytest.fixture(_isolation, scope="class")
     _function_isolation = pytest.fixture(_isolation, scope="function")
+
+    @allow_disconnected
+    def _snapshot(self) -> Optional[SnapshotID]:
+        try:
+            return self.chain_manager.snapshot()
+        except NotImplementedError:
+            if not self._warned_for_unimplemented_snapshot:
+                logger.warning(
+                    "The connected provider does not support snapshotting. "
+                    "Tests will not be completely isolated."
+                )
+                self._warned_for_unimplemented_snapshot = True
+
+        return None
+
+    @allow_disconnected
+    def _restore(self, snapshot_id: SnapshotID):
+        if snapshot_id not in self.chain_manager._snapshots:
+            return
+
+        self.chain_manager.restore(snapshot_id)
+
+
+class ReceiptCapture(ManagerAccessMixin):
+    pytest_config: PytestConfig
+    gas_report: Optional[GasReport] = None
+    receipt_map: Dict[str, ReceiptAPI] = {}
+    enter_blocks: List[int] = []
+
+    def __init__(self, pytest_config):
+        self.pytest_config = pytest_config
+
+    def __enter__(self):
+        block_number = self._get_block_number()
+        if block_number is not None:
+            self.enter_blocks.append(block_number)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if not self.enter_blocks:
+            return
+
+        start_block = self.enter_blocks.pop()
+        stop_block = self._get_block_number()
+        if stop_block is None or start_block > stop_block:
+            return
+
+        self.capture_range(start_block, stop_block)
+
+    @cached_property
+    def _track_gas(self) -> bool:
+        return self.pytest_config.getoption("--gas")
+
+    def capture_range(self, start_block: int, stop_block: int):
+        blocks = self.chain_manager.blocks.range(start_block, stop_block + 1)
+        transactions = [
+            t
+            for b in blocks
+            for t in b.transactions
+            if t.receiver and t.sender and t.sender in self.chain_manager.account_history
+        ]
+
+        for txn in transactions:
+            self.capture(txn.txn_hash.hex())
+
+    def capture(self, transaction_hash: str, track_gas: Optional[bool] = None):
+        if transaction_hash in self.receipt_map:
+            return
+
+        receipt = self.chain_manager.account_history.get_receipt(transaction_hash)
+        self.receipt_map[transaction_hash] = receipt
+        if not receipt:
+            return
+
+        if not receipt.receiver:
+            # TODO: Handle deploy receipts once trace supports it
+            return
+
+        # Merge-in the receipt's gas report with everything so far.
+        call_tree = receipt.call_tree
+        do_track_gas = track_gas if track_gas is not None else self._track_gas
+        if do_track_gas and call_tree:
+            parser = CallTraceParser(receipt)
+            gas_report = parser._get_rich_gas_report(call_tree)
+            if self.gas_report:
+                self.gas_report = merge_reports(self.gas_report, gas_report)
+            else:
+                self.gas_report = gas_report
+
+    @allow_disconnected
+    def _get_block_number(self) -> Optional[int]:
+        return self.provider.get_block("latest").number

src/ape/pytest/plugin.py

@@ -4,7 +4,7 @@
 import pytest
 
 from ape import networks, project
-from ape.pytest.fixtures import PytestApeFixtures
+from ape.pytest.fixtures import PytestApeFixtures, ReceiptCapture
 from ape.pytest.runners import PytestApeRunner
 
 
@@ -17,18 +17,23 @@ def pytest_addoption(parser):
         "--network",
         action="store",
         default=networks.default_ecosystem.name,
-        help="Override the default network and provider. (see ``ape networks list`` for options)",
+        help="Override the default network and provider (see ``ape networks list`` for options).",
     )
     parser.addoption(
         "--interactive",
         "-I",
         action="store_true",
-        help="Open an interactive console each time a test fails",
+        help="Open an interactive console each time a test fails.",
     )
     parser.addoption(
         "--disable-isolation",
         action="store_true",
-        help="Disable test and fixture isolation (see provider for info on snapshot availability)",
+        help="Disable test and fixture isolation (see provider for info on snapshot availability).",
+    )
+    parser.addoption(
+        "--gas",
+        action="store_true",
+        help="Show a transaction gas report at the end of the test session.",
     )
 
     # NOTE: Other pytest plugins, such as hypothesis, should integrate with pytest separately
@@ -46,19 +51,23 @@ def is_module(v):
         for module in modules:
             module.__tracebackhide__ = True
 
+    receipt_capture = ReceiptCapture(config)
+
     # Enable verbose output if stdout capture is disabled
     config.option.verbose = config.getoption("capture") == "no"
 
-    session = PytestApeRunner(pytest_config=config)
+    # Register the custom Ape test runner
+    session = PytestApeRunner(config, receipt_capture)
     config.pluginmanager.register(session, "ape-test")
 
-    fixtures = PytestApeFixtures()
+    # Include custom fixtures for project, accounts etc.
+    fixtures = PytestApeFixtures(config, receipt_capture)
     config.pluginmanager.register(fixtures, "ape-fixtures")
 
 
 def pytest_load_initial_conftests(early_config):
     """
-    Compile contracts before loading conftests.
+    Compile contracts before loading ``conftest.py``s.
     """
     cap_sys = early_config.pluginmanager.get_plugin("capturemanager")
 

src/ape/pytest/runners.py

@@ -3,21 +3,25 @@
 import click
 import pytest
 from _pytest.config import Config as PytestConfig
+from rich import print as rich_print
 
 import ape
 from ape.api import ProviderContextManager
-from ape.logging import logger
+from ape.logging import LogLevel, logger
 from ape.pytest.contextmanagers import RevertsContextManager
-from ape.utils import ManagerAccessMixin
+from ape.pytest.fixtures import ReceiptCapture
+from ape.utils import ManagerAccessMixin, parse_gas_table
 from ape_console._cli import console
 
 
 class PytestApeRunner(ManagerAccessMixin):
     def __init__(
         self,
         pytest_config: PytestConfig,
+        receipt_capture: ReceiptCapture,
     ):
         self.pytest_config = pytest_config
+        self.receipt_capture = receipt_capture
         self._provider_is_connected = False
         ape.reverts = RevertsContextManager  # type: ignore
 
@@ -38,7 +42,6 @@ def pytest_exception_interact(self, report, call):
         """
 
         if self.pytest_config.getoption("interactive") and report.failed:
-
             capman = self.pytest_config.pluginmanager.get_plugin("capturemanager")
             if capman:
                 capman.suspend_global_capture(in_=True)
@@ -153,14 +156,21 @@ def pytest_collection_finish(self, session):
             self._provider_context.push_provider()
             self._provider_is_connected = True
 
-    def pytest_sessionfinish(self):
+    def pytest_terminal_summary(self, terminalreporter):
         """
-        Called after whole test run finished, right before returning the exit
-        status to the system.
-
-        **NOTE**: This hook fires even when exceptions occur, so we cannot
-        assume the provider successfully connected.
+        Add a section to terminal summary reporting.
+        When ``--gas`` is active, outputs the gas profile report.
         """
+        if self.pytest_config.getoption("--gas"):
+            terminalreporter.section("Gas Profile")
+            gas_report = self.receipt_capture.gas_report
+            if gas_report:
+                tables = parse_gas_table(gas_report)
+                rich_print(*tables)
+            else:
+                terminalreporter.write_line(f"{LogLevel.WARNING.name}: No gas usage data found.")
+
+    def pytest_unconfigure(self):
         if self._provider_is_connected:
             self._provider_context.disconnect_all()
             self._provider_is_connected = False