From 585936f8be9ff915b5a480ae3e1e031c85610c7a Mon Sep 17 00:00:00 2001
From: slush <sam.loesch@gmail.com>
Date: Wed, 14 Aug 2024 17:43:58 -0500
Subject: [PATCH] fix: adds exception check for debug_setHead (#2217)

Co-authored-by: Juliya Smith <jules@apeworx.io>
---
 src/ape/pytest/fixtures.py        | 36 ++++++++++++++++++++-----------
 tests/functional/test_fixtures.py | 31 ++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/src/ape/pytest/fixtures.py b/src/ape/pytest/fixtures.py
index 0e5d179a45..71245166cf 100644
--- a/src/ape/pytest/fixtures.py
+++ b/src/ape/pytest/fixtures.py
@@ -23,7 +23,7 @@ class PytestApeFixtures(ManagerAccessMixin):
     # for fixtures, as they are used in output from the command
     # `ape test -q --fixture` (`pytest -q --fixture`).
 
-    _warned_for_unimplemented_snapshot = False
+    _supports_snapshot: bool = True
     receipt_capture: "ReceiptCapture"
 
     def __init__(self, config_wrapper: ConfigWrapper, receipt_capture: "ReceiptCapture"):
@@ -86,10 +86,13 @@ 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.
         """
-        try:
-            snapshot_id = self._snapshot()
-        except BlockNotFoundError:
-            snapshot_id = None
+        snapshot_id = None
+
+        if self._supports_snapshot:
+            try:
+                snapshot_id = self._snapshot()
+            except BlockNotFoundError:
+                self._supports_snapshot = False
 
         if self._track_transactions:
             did_yield = False
@@ -121,12 +124,12 @@ 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
+            logger.warning(
+                "The connected provider does not support snapshotting. "
+                "Tests will not be completely isolated."
+            )
+            # To avoid trying again
+            self._supports_snapshot = False
 
         return None
 
@@ -134,8 +137,15 @@ def _snapshot(self) -> Optional[SnapshotID]:
     def _restore(self, snapshot_id: SnapshotID):
         if snapshot_id not in self.chain_manager._snapshots:
             return
-
-        self.chain_manager.restore(snapshot_id)
+        try:
+            self.chain_manager.restore(snapshot_id)
+        except NotImplementedError:
+            logger.warning(
+                "The connected provider does not support snapshotting. "
+                "Tests will not be completely isolated."
+            )
+            # To avoid trying again
+            self._supports_snapshot = False
 
 
 class ReceiptCapture(ManagerAccessMixin):
diff --git a/tests/functional/test_fixtures.py b/tests/functional/test_fixtures.py
index dd9ee8919b..b27b3979f7 100644
--- a/tests/functional/test_fixtures.py
+++ b/tests/functional/test_fixtures.py
@@ -29,3 +29,34 @@ def test_isolation(isolation, receipt_capture):
     assert next(isolation) is None
     with pytest.raises(StopIteration):
         next(isolation)
+
+
+def test_isolation_restore_not_implemented(mocker, networks, fixtures):
+    isolation = fixtures._isolation()
+    mock_provider = mocker.MagicMock()
+    mock_provider.restore.side_effect = NotImplementedError
+    mock_provider.snapshot.return_value = 123
+    orig_provider = networks.active_provider
+    networks.active_provider = mock_provider
+    fixtures._supports_snapshot = True
+
+    try:
+        _ = next(isolation)
+        assert mock_provider.snapshot.call_count == 1
+        with pytest.raises(StopIteration):
+            _ = next(isolation)
+
+        # Is false because of the not-implemented error side-effect.
+        assert fixtures._supports_snapshot is False
+
+        isolation = fixtures._isolation()
+        _ = next(isolation)
+        # It does not call snapshot again.
+        assert mock_provider.snapshot.call_count == 1
+        with pytest.raises(StopIteration):
+            _ = next(isolation)
+
+        assert mock_provider.snapshot.call_count == 1
+
+    finally:
+        networks.active_provider = orig_provider