alexei.dolgolyov
eeb51fa4e7
Tests that call get_available_displays() or capture_display() require a real display ($DISPLAY on Linux, always available on Windows/macOS). Mark them with @requires_display to skip on headless CI instead of failing with "$DISPLAY not set". Affects: test_screen_capture.py (4 tests), test_api.py (1 test).
236 lines
7.1 KiB
Python
236 lines
7.1 KiB
Python
"""Tests for screen capture functionality."""
|
|
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from wled_controller.core.capture.screen_capture import (
|
|
get_available_displays,
|
|
capture_display,
|
|
extract_border_pixels,
|
|
get_edge_segments,
|
|
calculate_average_color,
|
|
calculate_median_color,
|
|
calculate_dominant_color,
|
|
ScreenCapture,
|
|
)
|
|
|
|
# Skip tests that require a real display on headless CI
|
|
_has_display = bool(os.environ.get("DISPLAY") or sys.platform == "win32" or sys.platform == "darwin")
|
|
requires_display = pytest.mark.skipif(not _has_display, reason="No display available (headless CI)")
|
|
|
|
|
|
@requires_display
|
|
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
|
|
|
|
|
|
@requires_display
|
|
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)
|
|
|
|
|
|
@requires_display
|
|
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")
|
|
|
|
# More segments than pixels is allowed — returns segments with fewer pixels each
|
|
result = get_edge_segments(edge_pixels, 200, "top")
|
|
assert len(result) == 200
|
|
|
|
|
|
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)
|
|
|
|
|
|
@requires_display
|
|
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)
|