chore: aligns the approach between snapshot and restore for catching …

…this error
ApeWorX · Aug 14, 2024 · f08bc9e · f08bc9e
1 parent f1e62d4
commit f08bc9e
Showing 1 changed file with 11 additions and 11 deletions.
diff --git a/src/ape/pytest/fixtures.py b/src/ape/pytest/fixtures.py
@@ -110,15 +110,7 @@ def _isolation(self) -> Iterator[None]:
             yield
 
         if snapshot_id is not None:
-            try:
-                self._restore(snapshot_id)
-            except NotImplementedError:
-                logger.warning(
-                    "The connected provider does not support snapshotting. "
-                    "Tests will not be completely isolated."
-                )
-                # Set to avoid trying again
-                self._supports_snapshot = False
+            self._restore(snapshot_id)
 
     # isolation fixtures
     _session_isolation = pytest.fixture(_isolation, scope="session")
@@ -136,6 +128,7 @@ def _snapshot(self) -> Optional[SnapshotID]:
                 "The connected provider does not support snapshotting. "
                 "Tests will not be completely isolated."
             )
+            # To avoid trying again
             self._supports_snapshot = False
 
         return None
@@ -144,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):