Simplify tests

Author: surajpaib

tests/unit/test_callbacks_writer_file.py

@@ -6,11 +6,12 @@
 import pytest
 import torch
 from PIL import Image
+import tempfile
 
 from lighter.callbacks.writer.file import LighterFileWriter
 
 
-def test_file_writer_initialization():
+def test_file_writer_initialization(tmp_path):
     """Test the initialization of LighterFileWriter class.
 
     This test verifies that:
@@ -21,17 +22,12 @@ def test_file_writer_initialization():
     The test creates a temporary directory, initializes a writer, checks its attributes,
     and then cleans up the directory.
     """
-    path = Path("test_dir")
-    path.mkdir(exist_ok=True)  # Ensure the directory exists
-    try:
-        writer = LighterFileWriter(path=path, writer="tensor")
-        assert writer.path == Path("test_dir")
-        assert writer.writer.__name__ == "write_tensor"  # Verify writer function
-    finally:
-        shutil.rmtree(path)  # Clean up after test
+    writer = LighterFileWriter(path=tmp_path, writer="tensor")
+    assert writer.path == tmp_path
+    assert writer.writer.__name__ == "write_tensor"  # Verify writer function
 
 
-def test_file_writer_write_tensor():
+def test_file_writer_write_tensor(tmp_path):
     """Test tensor writing functionality of LighterFileWriter.
 
     This test verifies that:
@@ -42,25 +38,20 @@ def test_file_writer_write_tensor():
     The test creates a simple tensor, saves it, loads it back, and verifies
     the content matches the original.
     """
-    test_dir = Path("test_dir")
-    test_dir.mkdir(exist_ok=True)
-    try:
-        writer = LighterFileWriter(path=test_dir, writer="tensor")
-        tensor = torch.tensor([1, 2, 3])
-        writer.write(tensor, id=1)
+    writer = LighterFileWriter(path=tmp_path, writer="tensor")
+    tensor = torch.tensor([1, 2, 3])
+    writer.write(tensor, id=1)
 
-        # Verify file exists
-        saved_path = writer.path / "1.pt"
-        assert saved_path.exists()
+    # Verify file exists
+    saved_path = writer.path / "1.pt"
+    assert saved_path.exists()
 
-        # Verify tensor contents
-        loaded_tensor = torch.load(saved_path)  # nosec B614
-        assert torch.equal(loaded_tensor, tensor)
-    finally:
-        shutil.rmtree(test_dir)
+    # Verify tensor contents
+    loaded_tensor = torch.load(saved_path)  # nosec B614
+    assert torch.equal(loaded_tensor, tensor)
 
 
-def test_file_writer_write_image():
+def test_file_writer_write_image(tmp_path):
     """Test image writing functionality of LighterFileWriter.
 
     This test verifies that:
@@ -71,26 +62,21 @@ def test_file_writer_write_image():
     The test creates a random RGB image tensor, saves it, and verifies
     the saved image properties.
     """
-    test_dir = Path("test_dir")
-    test_dir.mkdir(exist_ok=True)
-    try:
-        writer = LighterFileWriter(path=test_dir, writer="image")
-        tensor = torch.randint(0, 256, (3, 64, 64), dtype=torch.uint8)
-        writer.write(tensor, id="image_test")
-
-        # Verify file exists
-        saved_path = writer.path / "image_test.png"
-        assert saved_path.exists()
-
-        # Verify image contents
-        image = Image.open(saved_path)
-        image_array = np.array(image)
-        assert image_array.shape == (64, 64, 3)
-    finally:
-        shutil.rmtree(test_dir)
+    writer = LighterFileWriter(path=tmp_path, writer="image")
+    tensor = torch.randint(0, 256, (3, 64, 64), dtype=torch.uint8)
+    writer.write(tensor, id="image_test")
+
+    # Verify file exists
+    saved_path = writer.path / "image_test.png"
+    assert saved_path.exists()
+
+    # Verify image contents
+    image = Image.open(saved_path)
+    image_array = np.array(image)
+    assert image_array.shape == (64, 64, 3)
 
 
-def test_file_writer_write_video():
+def test_file_writer_write_video(tmp_path):
     """Test video writing functionality of LighterFileWriter.
 
     This test verifies that:
@@ -100,21 +86,16 @@ def test_file_writer_write_video():
     The test creates a random RGB video tensor and verifies it can be saved
     to disk in the correct format.
     """
-    test_dir = Path("test_dir")
-    test_dir.mkdir(exist_ok=True)
-    try:
-        writer = LighterFileWriter(path=test_dir, writer="video")
-        tensor = torch.randint(0, 256, (3, 10, 64, 64), dtype=torch.uint8)
-        writer.write(tensor, id="video_test")
+    writer = LighterFileWriter(path=tmp_path, writer="video")
+    tensor = torch.randint(0, 256, (3, 10, 64, 64), dtype=torch.uint8)
+    writer.write(tensor, id="video_test")
 
-        # Verify file exists
-        saved_path = writer.path / "video_test.mp4"
-        assert saved_path.exists()
-    finally:
-        shutil.rmtree(test_dir)
+    # Verify file exists
+    saved_path = writer.path / "video_test.mp4"
+    assert saved_path.exists()
 
 
-def test_file_writer_write_grayscale_video():
+def test_file_writer_write_grayscale_video(tmp_path):
     """Test grayscale video writing functionality of LighterFileWriter.
 
     This test verifies that:
@@ -125,22 +106,17 @@ def test_file_writer_write_grayscale_video():
     The test creates a grayscale video tensor and verifies it can be properly
     converted and saved as an MP4 file.
     """
-    test_dir = Path("test_dir")
-    test_dir.mkdir(exist_ok=True)
-    try:
-        writer = LighterFileWriter(path=test_dir, writer="video")
-        # Create a grayscale video tensor with 1 channel
-        tensor = torch.randint(0, 256, (1, 10, 64, 64), dtype=torch.uint8)
-        writer.write(tensor, id="grayscale_video_test")
-
-        # Verify file exists
-        saved_path = writer.path / "grayscale_video_test.mp4"
-        assert saved_path.exists()
-    finally:
-        shutil.rmtree(test_dir)
+    writer = LighterFileWriter(path=tmp_path, writer="video")
+    # Create a grayscale video tensor with 1 channel
+    tensor = torch.randint(0, 256, (1, 10, 64, 64), dtype=torch.uint8)
+    writer.write(tensor, id="grayscale_video_test")
 
+    # Verify file exists
+    saved_path = writer.path / "grayscale_video_test.mp4"
+    assert saved_path.exists()
 
-def test_file_writer_write_itk_image():
+
+def test_file_writer_write_itk_image(tmp_path):
     """Test ITK image writing functionality of LighterFileWriter.
 
     This test verifies that:
@@ -151,25 +127,20 @@ def test_file_writer_write_itk_image():
     The test attempts to write both regular tensors and MetaTensors,
     verifying proper error handling and successful writes.
     """
-    test_dir = Path("test_dir")
-    test_dir.mkdir(exist_ok=True)
-    try:
-        writer = LighterFileWriter(path=test_dir, writer="itk_nrrd")
-        tensor = torch.rand(1, 1, 64, 64, 64)  # Example 3D tensor
+    writer = LighterFileWriter(path=tmp_path, writer="itk_nrrd")
+    tensor = torch.rand(1, 1, 64, 64, 64)  # Example 3D tensor
 
-        # Test with regular tensor
-        with pytest.raises(TypeError, match="Tensor must be in MONAI MetaTensor format"):
-            writer.write(tensor, id="itk_image_test")
+    # Test with regular tensor
+    with pytest.raises(TypeError, match="Tensor must be in MONAI MetaTensor format"):
+        writer.write(tensor, id="itk_image_test")
 
-        # Test with proper MetaTensor
-        meta_tensor = monai.data.MetaTensor(tensor, affine=torch.eye(4), meta={"original_channel_dim": 1})
-        writer.write(meta_tensor, id="itk_image_test")
+    # Test with proper MetaTensor
+    meta_tensor = monai.data.MetaTensor(tensor, affine=torch.eye(4), meta={"original_channel_dim": 1})
+    writer.write(meta_tensor, id="itk_image_test")
 
-        # Verify file exists
-        saved_path = writer.path / "itk_image_test.nrrd"
-        assert saved_path.exists()
-    finally:
-        shutil.rmtree(test_dir)
+    # Verify file exists
+    saved_path = writer.path / "itk_image_test.nrrd"
+    assert saved_path.exists()
 
 
 def test_file_writer_invalid_directory():