wled-screen-controller-mixed/server/tests/test_screen_capture.py

"""Tests for screen capture functionality."""

import numpy as np
import pytest

from wled_controller.core.screen_capture import (
    get_available_displays,
    capture_display,
    extract_border_pixels,
    get_edge_segments,
    calculate_average_color,
    calculate_median_color,
    calculate_dominant_color,
    ScreenCapture,
)


def test_get_available_displays():
    """Test getting available displays."""
    displays = get_available_displays()

    assert isinstance(displays, list)
    assert len(displays) >= 1  # At least one display should be available

    # Check first display structure
    display = displays[0]
    assert hasattr(display, "index")
    assert hasattr(display, "name")
    assert hasattr(display, "width")
    assert hasattr(display, "height")
    assert display.width > 0
    assert display.height > 0


def test_capture_display():
    """Test capturing a display."""
    # Capture the first display
    capture = capture_display(0)

    assert isinstance(capture, ScreenCapture)
    assert capture.image is not None
    assert capture.width > 0
    assert capture.height > 0
    assert capture.display_index == 0
    assert isinstance(capture.image, np.ndarray)
    assert capture.image.shape == (capture.height, capture.width, 3)


def test_capture_display_invalid_index():
    """Test capturing with invalid display index."""
    with pytest.raises(ValueError):
        capture_display(999)  # Invalid display index


def test_extract_border_pixels():
    """Test extracting border pixels."""
    # Create a test screen capture
    test_image = np.random.randint(0, 256, (100, 200, 3), dtype=np.uint8)
    capture = ScreenCapture(
        image=test_image,
        width=200,
        height=100,
        display_index=0
    )

    border_width = 10
    borders = extract_border_pixels(capture, border_width)

    # Check border shapes
    assert borders.top.shape == (border_width, 200, 3)
    assert borders.bottom.shape == (border_width, 200, 3)
    assert borders.left.shape == (100, border_width, 3)
    assert borders.right.shape == (100, border_width, 3)


def test_extract_border_pixels_invalid_width():
    """Test extracting borders with invalid width."""
    test_image = np.random.randint(0, 256, (100, 200, 3), dtype=np.uint8)
    capture = ScreenCapture(
        image=test_image,
        width=200,
        height=100,
        display_index=0
    )

    # Border width too small
    with pytest.raises(ValueError):
        extract_border_pixels(capture, 0)

    # Border width too large
    with pytest.raises(ValueError):
        extract_border_pixels(capture, 50)


def test_get_edge_segments():
    """Test dividing edge into segments."""
    # Create test edge pixels (horizontal edge)
    edge_pixels = np.random.randint(0, 256, (10, 100, 3), dtype=np.uint8)

    segments = get_edge_segments(edge_pixels, 10, "top")

    assert len(segments) == 10
    # Each segment should have width of approximately 10
    for segment in segments:
        assert segment.shape[0] == 10  # Height stays same
        assert 8 <= segment.shape[1] <= 12  # Width varies slightly
        assert segment.shape[2] == 3  # RGB


def test_get_edge_segments_vertical():
    """Test dividing vertical edge into segments."""
    # Create test edge pixels (vertical edge)
    edge_pixels = np.random.randint(0, 256, (100, 10, 3), dtype=np.uint8)

    segments = get_edge_segments(edge_pixels, 10, "left")

    assert len(segments) == 10
    # Each segment should have height of approximately 10
    for segment in segments:
        assert 8 <= segment.shape[0] <= 12  # Height varies slightly
        assert segment.shape[1] == 10  # Width stays same
        assert segment.shape[2] == 3  # RGB


def test_get_edge_segments_invalid():
    """Test edge segments with invalid parameters."""
    edge_pixels = np.random.randint(0, 256, (10, 100, 3), dtype=np.uint8)

    with pytest.raises(ValueError):
        get_edge_segments(edge_pixels, 0, "top")

    with pytest.raises(ValueError):
        get_edge_segments(edge_pixels, 200, "top")  # More segments than pixels


def test_calculate_average_color():
    """Test calculating average color."""
    # Create uniform color region
    pixels = np.full((10, 10, 3), [100, 150, 200], dtype=np.uint8)

    color = calculate_average_color(pixels)

    assert color == (100, 150, 200)


def test_calculate_average_color_mixed():
    """Test average color with mixed colors."""
    # Create region with two colors
    pixels = np.zeros((10, 10, 3), dtype=np.uint8)
    pixels[:5, :, :] = [255, 0, 0]  # Top half red
    pixels[5:, :, :] = [0, 0, 255]  # Bottom half blue

    color = calculate_average_color(pixels)

    # Should be roughly purple (average of red and blue)
    assert 120 <= color[0] <= 135  # R
    assert 0 <= color[1] <= 10  # G
    assert 120 <= color[2] <= 135  # B


def test_calculate_median_color():
    """Test calculating median color."""
    # Create region with outliers
    pixels = np.full((10, 10, 3), [100, 100, 100], dtype=np.uint8)
    pixels[0, 0, :] = [255, 255, 255]  # One bright outlier

    color = calculate_median_color(pixels)

    # Median should be close to 100, not affected by outlier
    assert 95 <= color[0] <= 105
    assert 95 <= color[1] <= 105
    assert 95 <= color[2] <= 105


def test_calculate_dominant_color():
    """Test calculating dominant color."""
    # Create region with mostly one color
    pixels = np.full((20, 20, 3), [100, 150, 200], dtype=np.uint8)
    # Add some noise
    pixels[:2, :2, :] = [50, 75, 100]

    color = calculate_dominant_color(pixels)

    # Dominant color should be close to the main color
    assert 90 <= color[0] <= 110
    assert 140 <= color[1] <= 160
    assert 190 <= color[2] <= 210


def test_calculate_color_empty_pixels():
    """Test color calculation with empty pixel array."""
    empty_pixels = np.array([]).reshape(0, 0, 3)

    assert calculate_average_color(empty_pixels) == (0, 0, 0)
    assert calculate_median_color(empty_pixels) == (0, 0, 0)
    assert calculate_dominant_color(empty_pixels) == (0, 0, 0)


def test_end_to_end_screen_capture():
    """Test complete screen capture workflow."""
    # Get available displays
    displays = get_available_displays()
    assert len(displays) > 0

    # Capture first display
    capture = capture_display(0)
    assert capture is not None

    # Extract borders
    borders = extract_border_pixels(capture, 10)
    assert borders.top is not None
    assert borders.bottom is not None
    assert borders.left is not None
    assert borders.right is not None

    # Get segments for top edge
    top_segments = get_edge_segments(borders.top, 10, "top")
    assert len(top_segments) == 10

    # Calculate color for first segment
    color = calculate_average_color(top_segments[0])
    assert len(color) == 3
    assert all(0 <= c <= 255 for c in color)