SNOW-892716 Improve localfs put/get test (#1786)

* SNOW-892716 Improve localfs put/get test

Description

This is a long overdue improvement from early slack message from
Brandon

Testing

unit test

* comment

Author: sfc-gh-sfan

test/unit/test_local_storage_client.py

@@ -2,11 +2,10 @@
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
 #
 
-import os
 import random
-import shutil
 import string
 import tempfile
+from pathlib import Path
 
 import pytest
 
@@ -26,18 +25,14 @@ def test_multi_chunk_upload(multipart_threshold):
         [random.choice(string.ascii_letters) for _ in range(300)]
     ).encode()
     file_name = "test_file"
-    stage_dir = tempfile.mkdtemp()
-    stage_file = os.path.join(stage_dir, file_name)
-    local_dir = tempfile.mkdtemp()
-    local_file = os.path.join(local_dir, file_name)
-
-    try:
-        with open(local_file, "wb+") as fd:
-            fd.write(file_content)
+    with tempfile.TemporaryDirectory() as stage_dir, tempfile.TemporaryDirectory() as local_dir:
+        stage_file = Path(stage_dir) / file_name
+        local_file = Path(local_dir) / file_name
+        Path(local_file).write_bytes(file_content)
 
         meta = SnowflakeFileMeta(
             name=file_name,
-            src_file_name=local_file,
+            src_file_name=str(local_file),
             stage_location_type=LOCAL_FS,
             dst_file_name=file_name,
             multipart_threshold=multipart_threshold,
@@ -47,11 +42,7 @@ def test_multi_chunk_upload(multipart_threshold):
         for chunk_id in range(client.num_of_chunks):
             client.upload_chunk(chunk_id)
 
-        with open(stage_file, "rb") as fd:
-            assert fd.read() == file_content
-    finally:
-        shutil.rmtree(stage_dir, ignore_errors=True)
-        shutil.rmtree(local_dir, ignore_errors=True)
+        assert Path(stage_file).read_bytes() == file_content
 
 
 @pytest.mark.parametrize("multipart_threshold", [0, 67108864])
@@ -60,18 +51,14 @@ def test_multi_chunk_download(multipart_threshold):
         [random.choice(string.ascii_letters) for _ in range(300)]
     ).encode()
     file_name = "test_file"
-    stage_dir = tempfile.mkdtemp()
-    stage_file = os.path.join(stage_dir, file_name)
-    local_dir = tempfile.mkdtemp()
-    local_file = os.path.join(local_dir, file_name)
-
-    try:
-        with open(stage_file, "wb+") as fd:
-            fd.write(file_content)
+    with tempfile.TemporaryDirectory() as stage_dir, tempfile.TemporaryDirectory() as local_dir:
+        stage_file = Path(stage_dir) / file_name
+        local_file = Path(local_dir) / file_name
+        Path(stage_file).write_bytes(file_content)
 
         meta = SnowflakeFileMeta(
             name=file_name,
-            src_file_name=stage_file,
+            src_file_name=str(stage_file),
             stage_location_type=LOCAL_FS,
             dst_file_name=file_name,
             local_location=local_dir,
@@ -83,8 +70,4 @@ def test_multi_chunk_download(multipart_threshold):
             client.download_chunk(chunk_id)
         client.finish_download()
 
-        with open(local_file, "rb") as fd:
-            assert fd.read() == file_content
-    finally:
-        shutil.rmtree(stage_dir, ignore_errors=True)
-        shutil.rmtree(local_dir, ignore_errors=True)
+        assert Path(local_file).read_bytes() == file_content