From ff3754e07e21aca9a128fff2a75d33475e001081 Mon Sep 17 00:00:00 2001
From: antazoey <jules@apeworx.io>
Date: Fri, 14 Jun 2024 10:46:08 -0500
Subject: [PATCH] feat: create `config.isolate_data_folder()` method (#2139)

---
 src/ape/managers/config.py | 25 ++++++++++++++++++++++++-
 tests/conftest.py          |  7 ++++---
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/ape/managers/config.py b/src/ape/managers/config.py
index 0f725fc46a..60aa8672a2 100644
--- a/src/ape/managers/config.py
+++ b/src/ape/managers/config.py
@@ -1,4 +1,6 @@
 import os
+from collections.abc import Iterator
+from contextlib import contextmanager
 from functools import cached_property
 from pathlib import Path
 from typing import Any, Optional
@@ -8,7 +10,7 @@
 from ape.api import PluginConfig
 from ape.api.config import ApeConfig
 from ape.managers.base import BaseManager
-from ape.utils import log_instead_of_fail
+from ape.utils import create_tempdir, in_tempdir, log_instead_of_fail
 from ape.utils.basemodel import (
     ExtraAttributesMixin,
     ExtraModelAttributes,
@@ -120,6 +122,27 @@ def get_config(self, plugin_name: str) -> PluginConfig:
             # Empty config.
             return PluginConfig()
 
+    @contextmanager
+    def isolate_data_folder(self) -> Iterator[Path]:
+        """
+        Change Ape's DATA_FOLDER to point a temporary path,
+        in a context, for testing purposes. Any data
+        cached to disk will not persist.
+        """
+        original_data_folder = self.DATA_FOLDER
+        if in_tempdir(original_data_folder):
+            # Already isolated.
+            yield original_data_folder
+
+        else:
+            try:
+                with create_tempdir() as temp_data_folder:
+                    self.DATA_FOLDER = temp_data_folder
+                    yield temp_data_folder
+
+            finally:
+                self.DATA_FOLDER = original_data_folder
+
 
 def merge_configs(*cfgs: dict) -> dict:
     if len(cfgs) == 0:
diff --git a/tests/conftest.py b/tests/conftest.py
index 980f924f65..67a809807a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,7 +3,6 @@
 import shutil
 import subprocess
 import sys
-import tempfile
 import time
 from collections.abc import Callable, Sequence
 from contextlib import contextmanager
@@ -30,7 +29,9 @@
 explorer_test = pytest.mark.xdist_group(name="explorer-tests")
 
 # Ensure we don't persist any .ape data or using existing.
-DATA_FOLDER = Path(tempfile.mkdtemp()).resolve()
+
+_DATA_FOLDER_CTX = ape.config.isolate_data_folder()
+DATA_FOLDER = _DATA_FOLDER_CTX.__enter__()
 ape.config.DATA_FOLDER = DATA_FOLDER
 
 
@@ -59,7 +60,7 @@ def setenviron(monkeypatch):
 @pytest.fixture(scope="session", autouse=True)
 def clean_temp_data_folder():
     yield
-    shutil.rmtree(DATA_FOLDER)
+    _DATA_FOLDER_CTX.__exit__(None, None, None)
 
 
 @pytest.fixture(scope="session", autouse=True)