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