"""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)