tests/test_visualization.py

"""Tests for the visualization module."""

import pytest
from PIL import Image

from src.parsing import BBox
from src.visualization import create_no_detection_overlay, draw_boxes


@pytest.fixture
def sample_image():
    return Image.new("RGB", (640, 480), color=(30, 30, 60))


@pytest.fixture
def sample_boxes():
    return [
        BBox(x1=50, y1=50, x2=200, y2=150, confidence=0.9, label="debris"),
        BBox(x1=300, y1=200, x2=450, y2=350, confidence=0.75, label="satellite"),
    ]


class TestDrawBoxes:
    def test_returns_image(self, sample_image, sample_boxes):
        result = draw_boxes(sample_image, sample_boxes)
        assert isinstance(result, Image.Image)
        assert result.size == sample_image.size

    def test_empty_boxes_returns_copy(self, sample_image):
        result = draw_boxes(sample_image, [])
        assert isinstance(result, Image.Image)
        assert result.size == sample_image.size

    def test_custom_labels(self, sample_image, sample_boxes):
        labels = ["fragment", "panel"]
        result = draw_boxes(sample_image, sample_boxes, labels=labels)
        assert isinstance(result, Image.Image)

    def test_tiny_boxes_expanded(self, sample_image):
        tiny_boxes = [BBox(x1=100, y1=100, x2=101, y2=101)]
        result = draw_boxes(sample_image, tiny_boxes)
        assert isinstance(result, Image.Image)

    def test_out_of_bounds_boxes_clipped(self, sample_image):
        boxes = [BBox(x1=-50, y1=-50, x2=800, y2=600)]
        result = draw_boxes(sample_image, boxes)
        assert isinstance(result, Image.Image)

    def test_no_confidence_display(self, sample_image, sample_boxes):
        result = draw_boxes(sample_image, sample_boxes, show_confidence=False)
        assert isinstance(result, Image.Image)


class TestNoDetectionOverlay:
    def test_returns_image(self, sample_image):
        result = create_no_detection_overlay(sample_image)
        assert isinstance(result, Image.Image)
        assert result.size == sample_image.size

    def test_custom_message(self, sample_image):
        result = create_no_detection_overlay(sample_image, "Custom message")
        assert isinstance(result, Image.Image)