diff --git a/src/prefect/testing/utilities.py b/src/prefect/testing/utilities.py index a6a187dbf5d3..29dd1e28b4d1 100644 --- a/src/prefect/testing/utilities.py +++ b/src/prefect/testing/utilities.py @@ -2,11 +2,13 @@ Internal utilities for tests. """ +import atexit +import shutil import warnings from contextlib import ExitStack, contextmanager from pathlib import Path from pprint import pprint -from tempfile import TemporaryDirectory +from tempfile import mkdtemp from typing import TYPE_CHECKING, Dict, List, Union import prefect.context @@ -118,24 +120,30 @@ def prefect_test_harness(): from prefect.server.database.dependencies import temporary_database_interface # create temp directory for the testing database - with TemporaryDirectory() as temp_dir: - with ExitStack() as stack: - # temporarily override any database interface components - stack.enter_context(temporary_database_interface()) - - DB_PATH = "sqlite+aiosqlite:///" + str(Path(temp_dir) / "prefect-test.db") - stack.enter_context( - prefect.settings.temporary_settings( - # Clear the PREFECT_API_URL - restore_defaults={prefect.settings.PREFECT_API_URL}, - # Use a temporary directory for the database - updates={ - prefect.settings.PREFECT_API_DATABASE_CONNECTION_URL: DB_PATH, - prefect.settings.PREFECT_API_URL: None, - }, - ) + temp_dir = mkdtemp() + + def cleanup_temp_dir(temp_dir): + shutil.rmtree(temp_dir) + + atexit.register(cleanup_temp_dir, temp_dir) + + with ExitStack() as stack: + # temporarily override any database interface components + stack.enter_context(temporary_database_interface()) + + DB_PATH = "sqlite+aiosqlite:///" + str(Path(temp_dir) / "prefect-test.db") + stack.enter_context( + prefect.settings.temporary_settings( + # Clear the PREFECT_API_URL + restore_defaults={prefect.settings.PREFECT_API_URL}, + # Use a temporary directory for the database + updates={ + prefect.settings.PREFECT_API_DATABASE_CONNECTION_URL: DB_PATH, + prefect.settings.PREFECT_API_URL: None, + }, ) - yield + ) + yield async def get_most_recent_flow_run(client: "PrefectClient" = None):