scan4llm/tests/test_slicer.py

from __future__ import annotations

import tempfile
from pathlib import Path

import cv2
import numpy as np
import pytest

from src.split.slicer import (
    SUPPORTED_EXTENSIONS,
    VALID_ROTATIONS,
    add_border,
    generate_output_paths,
    load_image,
    parse_slice,
    process_image,
    rotate_image,
    save_page,
    slice_grid,
)
from tests.helpers import create_test_image, image_to_path


class TestParseSlice:
    def test_simple_2x2(self) -> None:
        assert parse_slice("2:2") == (2, 2)

    def test_2x4(self) -> None:
        assert parse_slice("2:4") == (2, 4)

    def test_1x1(self) -> None:
        assert parse_slice("1:1") == (1, 1)

    def test_large(self) -> None:
        assert parse_slice("10:10") == (10, 10)

    def test_invalid_format_missing_colon(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("22")

    def test_invalid_format_extra_colon(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("2:3:4")

    def test_non_integer_value(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("a:b")

    def test_non_integer_col(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("x:2")

    def test_zero_cols(self) -> None:
        with pytest.raises(ValueError, match="Неверное значение --slice"):
            parse_slice("0:2")

    def test_zero_rows(self) -> None:
        with pytest.raises(ValueError, match="Неверное значение --slice"):
            parse_slice("2:0")

    def test_negative_cols(self) -> None:
        with pytest.raises(ValueError, match="Неверное значение --slice"):
            parse_slice("-1:2")

    def test_blank_string(self) -> None:
        with pytest.raises(ValueError):
            parse_slice("")


class TestSliceGrid:
    def test_2x2_exact(self) -> None:
        image = create_test_image(400, 300)
        pages = slice_grid(image, 2, 2)
        assert len(pages) == 4
        assert pages[0].shape == (150, 200, 3)
        assert pages[1].shape == (150, 200, 3)
        assert pages[2].shape == (150, 200, 3)
        assert pages[3].shape == (150, 200, 3)

    def test_1x1(self) -> None:
        image = create_test_image(100, 100)
        pages = slice_grid(image, 1, 1)
        assert len(pages) == 1
        assert pages[0].shape == (100, 100, 3)

    def test_2x3(self) -> None:
        image = create_test_image(300, 200)
        pages = slice_grid(image, 2, 3)
        assert len(pages) == 6
        for page in pages:
            assert page.shape == (100, 100, 3)

    def test_uneven_division_width(self) -> None:
        image = create_test_image(width=403, height=300)
        pages = slice_grid(image, 1, 2)
        assert len(pages) == 2
        assert pages[0].shape == (300, 201, 3)
        assert pages[1].shape == (300, 202, 3)

    def test_uneven_division_height(self) -> None:
        image = create_test_image(width=400, height=301)
        pages = slice_grid(image, 2, 1)
        assert len(pages) == 2
        assert pages[0].shape == (150, 400, 3)
        assert pages[1].shape == (151, 400, 3)

    def test_uneven_both(self) -> None:
        image = create_test_image(width=401, height=301)
        pages = slice_grid(image, 2, 2)
        assert len(pages) == 4
        assert pages[0].shape == (150, 200, 3)
        assert pages[1].shape == (150, 201, 3)
        assert pages[2].shape == (151, 200, 3)
        assert pages[3].shape == (151, 201, 3)

    def test_pixel_content_preserved(self) -> None:
        image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
        pages = slice_grid(image, 2, 2)
        assert np.array_equal(image[0:50, 0:50], pages[0])
        assert np.array_equal(image[0:50, 50:100], pages[1])
        assert np.array_equal(image[50:100, 0:50], pages[2])
        assert np.array_equal(image[50:100, 50:100], pages[3])

    def test_total_pixels_preserved(self) -> None:
        image = create_test_image(width=401, height=301)
        pages = slice_grid(image, 3, 3)
        total = sum(page.shape[0] * page.shape[1] for page in pages)
        assert total == 401 * 301

    def test_order_2x2(self) -> None:
        image = create_test_image(200, 200)
        image[0:100, 0:100] = (255, 0, 0)
        image[0:100, 100:200] = (0, 255, 0)
        image[100:200, 0:100] = (0, 0, 255)
        image[100:200, 100:200] = (255, 255, 0)
        pages = slice_grid(image, 2, 2)
        assert pages[0][0, 0].tolist() == [255, 0, 0]
        assert pages[1][0, 0].tolist() == [0, 255, 0]
        assert pages[2][0, 0].tolist() == [0, 0, 255]
        assert pages[3][0, 0].tolist() == [255, 255, 0]

    def test_order_2x4(self) -> None:
        image = create_test_image(400, 200)
        pages = slice_grid(image, 2, 4)
        assert len(pages) == 8
        assert pages[0].shape == (100, 100, 3)
        assert pages[7].shape == (100, 100, 3)

    def test_4x4(self) -> None:
        image = create_test_image(400, 400)
        pages = slice_grid(image, 4, 4)
        assert len(pages) == 16
        for page in pages:
            assert page.shape == (100, 100, 3)


class TestAddBorder:
    def test_default_border(self) -> None:
        image = create_test_image(100, 100)
        result = add_border(image, 20)
        assert result.shape == (140, 140, 3)

    def test_zero_border(self) -> None:
        image = create_test_image(100, 100)
        result = add_border(image, 0)
        assert np.array_equal(result, image)

    def test_border_is_white(self) -> None:
        image = create_test_image(100, 100)
        border = 10
        result = add_border(image, border)
        assert np.all(result[0:border, :] == 255)
        assert np.all(result[-border:, :] == 255)
        assert np.all(result[:, 0:border] == 255)
        assert np.all(result[:, -border:] == 255)

    def test_large_border(self) -> None:
        image = create_test_image(10, 10)
        result = add_border(image, 50)
        assert result.shape == (110, 110, 3)


class TestRotateImage:
    def test_rotate_90(self) -> None:
        image = create_test_image(200, 100)
        result = rotate_image(image, 90)
        assert result.shape == (200, 100, 3)

    def test_rotate_180(self) -> None:
        image = create_test_image(200, 100)
        result = rotate_image(image, 180)
        assert result.shape == (100, 200, 3)

    def test_rotate_270(self) -> None:
        image = create_test_image(200, 100)
        result = rotate_image(image, 270)
        assert result.shape == (200, 100, 3)

    def test_invalid_angle(self) -> None:
        image = create_test_image()
        with pytest.raises(ValueError, match="Недопустимый угол поворота"):
            rotate_image(image, 45)

    def test_invalid_angle_zero(self) -> None:
        image = create_test_image()
        with pytest.raises(ValueError):
            rotate_image(image, 0)

    def test_valid_rotations_set(self) -> None:
        assert VALID_ROTATIONS == frozenset({90, 180, 270})


class TestGenerateOutputPaths:
    def test_simple(self) -> None:
        paths = generate_output_paths(Path("scan001.jpg"), 2, Path("/tmp"))
        assert paths == [Path("/tmp/scan001_01.jpg"), Path("/tmp/scan001_02.jpg")]

    def test_tiff_extension(self) -> None:
        paths = generate_output_paths(Path("doc.tiff"), 3, Path("out"))
        assert paths == [
            Path("out/doc_01.tiff"),
            Path("out/doc_02.tiff"),
            Path("out/doc_03.tiff"),
        ]

    def test_png_extension(self) -> None:
        paths = generate_output_paths(Path("img.PNG"), 1, Path("out"))
        assert paths == [Path("out/img_01.png")]

    def test_zero_padding(self) -> None:
        paths = generate_output_paths(Path("scan.jpg"), 12, Path("out"))
        assert paths[0] == Path("out/scan_01.jpg")
        assert paths[9] == Path("out/scan_10.jpg")
        assert paths[11] == Path("out/scan_12.jpg")

    def test_many_pages(self) -> None:
        paths = generate_output_paths(Path("scan.jpg"), 100, Path("out"))
        assert paths[0] == Path("out/scan_001.jpg")
        assert paths[99] == Path("out/scan_100.jpg")

    def test_dot_in_filename(self) -> None:
        paths = generate_output_paths(Path("scan.001.jpg"), 2, Path("out"))
        assert paths == [Path("out/scan.001_01.jpg"), Path("out/scan.001_02.jpg")]


class TestLoadImage:
    def test_valid_jpeg(self) -> None:
        image = create_test_image()
        path = image_to_path(image, ".jpg")
        loaded = load_image(path)
        assert loaded.shape == (300, 400, 3)

    def test_valid_png(self) -> None:
        image = create_test_image()
        path = image_to_path(image, ".png")
        loaded = load_image(path)
        assert loaded.shape == (300, 400, 3)

    def test_valid_tiff(self) -> None:
        image = create_test_image()
        path = image_to_path(image, ".tiff")
        loaded = load_image(path)
        assert loaded.shape == (300, 400, 3)

    def test_file_not_found(self) -> None:
        with pytest.raises(FileNotFoundError, match="Файл не найден"):
            load_image(Path("/tmp/nonexistent_scan2html_test.jpg"))

    def test_unsupported_extension(self) -> None:
        path = Path("/tmp/test.bmp")
        with pytest.raises(ValueError, match="Неподдерживаемый формат"):
            load_image(path)

    def test_unsupported_extension_no_dot(self) -> None:
        path = Path("/tmp/test")
        with pytest.raises(ValueError, match="Неподдерживаемый формат"):
            load_image(path)

    def test_supported_extensions_set(self) -> None:
        assert SUPPORTED_EXTENSIONS == frozenset({".jpg", ".jpeg", ".png", ".tiff", ".tif"})


class TestSavePage:
    def test_saves_file(self) -> None:
        image = create_test_image()
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "test.png"
            save_page(image, path)
            assert path.exists()
            loaded = cv2.imread(str(path))
            assert loaded is not None
            assert loaded.shape == (300, 400, 3)

    def test_creates_parent_dir(self) -> None:
        image = create_test_image()
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "subdir" / "test.png"
            save_page(image, path)
            assert path.exists()


class TestProcessImage:
    def test_end_to_end_2x2(self) -> None:
        image = create_test_image(400, 300)
        path = image_to_path(image, ".jpg")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp))
            assert len(result) == 4
            for p in result:
                assert p.exists()
                loaded = cv2.imread(str(p))
                assert loaded.shape == (250, 300, 3)

    def test_end_to_end_with_border(self) -> None:
        image = create_test_image(200, 200)
        path = image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp), border_px=10)
            assert len(result) == 4
            for p in result:
                assert p.exists()
                loaded = cv2.imread(str(p))
                assert loaded.shape == (120, 120, 3)

    def test_end_to_end_tiff(self) -> None:
        image = create_test_image(400, 400)
        path = image_to_path(image, ".tiff")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp))
            assert len(result) == 4
            for p in result:
                assert p.suffix.lower() == ".tiff"
                assert p.exists()

    def test_end_to_end_with_rotation(self) -> None:
        image = create_test_image(200, 400)
        path = image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp), pre_rotate=90)
            assert len(result) == 4
            for p in result:
                assert p.exists()
            loaded = cv2.imread(str(result[0]))
            assert loaded.shape[0] > 0