From 7f09f2f3ff6193c4c014f3b74ececc8a6845e40d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=9F=E5=AE=B6=E7=91=8B?= <36886416+JiangJiaWei1103@users.noreply.github.com> Date: Tue, 12 Nov 2024 14:10:05 +0100 Subject: [PATCH] test: Use context manager for tmp dirs (#2922) Signed-off-by: JiaWei Jiang Signed-off-by: 400Ping <43886578+400Ping@users.noreply.github.com> --- .../flytekit/unit/types/directory/test_dir.py | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/tests/flytekit/unit/types/directory/test_dir.py b/tests/flytekit/unit/types/directory/test_dir.py index 9d311cedae..f4d7668dba 100644 --- a/tests/flytekit/unit/types/directory/test_dir.py +++ b/tests/flytekit/unit/types/directory/test_dir.py @@ -12,20 +12,14 @@ @pytest.fixture def local_tmp_dirs(): - # Create a source directory - src_dir = tempfile.TemporaryDirectory() - for file_idx in range(N_FILES_PER_DIR): - with open(Path(src_dir.name) / f"{file_idx}.txt", "w") as f: - f.write(str(file_idx)) - - # Create an empty directory as the destination - dst_dir = tempfile.TemporaryDirectory() - - yield src_dir.name, dst_dir.name + # Create a source and an empty destination directory + with (tempfile.TemporaryDirectory() as src_dir, + tempfile.TemporaryDirectory() as dst_dir): + for file_idx in range(N_FILES_PER_DIR): + with open(Path(src_dir) / f"{file_idx}.txt", "w") as f: + f.write(str(file_idx)) - # Cleanup - src_dir.cleanup() - dst_dir.cleanup() + yield src_dir, dst_dir def test_src_path_with_different_types(local_tmp_dirs) -> None: