Merge pull request #9 from este6an13/test

checks-ocr\test(back): #5 add pdf to img unit test

src/utils/__init__.py

@@ -5,7 +5,7 @@
 from PIL import Image
 
 
-def calculate_iou(box1, box2):
+def calculate_iou(box1: tuple[float], box2: tuple[float]) -> float:
     """
     Calculate Intersection over Union (IoU) between two bounding boxes.
 
@@ -34,8 +34,8 @@ def calculate_iou(box1, box2):
 
     return iou
 
-def pdf_to_img(pdf_path, folder, resolution=300):
-    file_path = os.path.join(folder, f"{pdf_path}")
+def pdf_to_img(file_name: str, folder: str, resolution: int = 300):
+    file_path = os.path.join(folder, file_name)
     file_handle = fitz.open(file_path)
 
     # Get the first page
@@ -55,8 +55,8 @@ def pdf_to_img(pdf_path, folder, resolution=300):
 
 
 def generate_id():
-  id = str(uuid.uuid4())
-  return id
+    id = str(uuid.uuid4())
+    return id
 
 def contains_number(s):
     """

test/unit/src/test_utils.py

@@ -1,6 +1,12 @@
-from src.utils import calculate_iou
-
+from src.utils import (
+    calculate_iou,
+    pdf_to_img,
+)
+from PIL import Image
+import fitz
+import os
 import pytest
+import shutil
 
 @pytest.mark.parametrize(
     ("box1", "box2", "expected"),
@@ -10,7 +16,47 @@
         ((0.2, 0.2, 0.0, 0.0), (0.2, 0.2, 0.1, 0.0), 1/3,),
     ],
 )
-def test_calculate_iou(box1: tuple, box2: tuple, expected: float) -> None:
+def test_calculate_iou(box1: tuple[float], box2: tuple[float], expected: float) -> None:
     iou = calculate_iou(box1, box2)
     assert iou >= 0 and iou <= 1
-    assert pytest.approx(iou) == expected
+    assert pytest.approx(iou) == expected
+
+
+def _create_sample_pdf(pdf_path: str) -> None:
+    doc = fitz.open()
+    doc.insert_page(0)
+    doc.save(str(pdf_path))
+
+
+@pytest.mark.parametrize(
+    ("resolution", "expected_width", "expected_height", "expected_format"),
+    [
+        (300, 2480, 3508, "RGB"),
+        (150, 1240, 1754, "RGB"),
+    ],
+)
+def test_pdf_to_img(resolution: int, expected_width: int, expected_height: int, expected_format: str) -> None:
+    folder = r"./test/unit/temp/"
+
+    # Create the output folder if it doesn't exist
+    os.makedirs(folder, exist_ok=True)
+
+    file_name = "sample.pdf"
+
+    pdf_path = os.path.join(folder, file_name)
+
+    _create_sample_pdf(pdf_path)
+
+    img = pdf_to_img(file_name, folder, resolution)
+
+    # Check if the output image is an instance of PIL Image
+    assert isinstance(img, Image.Image)
+
+    # Check image dimensions
+    assert pytest.approx(img.size, rel=1.0) == (expected_width, expected_height)
+
+    # Check image format
+    assert img.mode == expected_format
+
+    # Clean up: Remove the temporary folder and its contents
+    shutil.rmtree(folder)