corrscope/tests/test_renderer.py

import numpy as np
import pytest
from matplotlib.colors import to_rgb

from corrscope.channel import ChannelConfig
from corrscope.outputs import RGB_DEPTH
from corrscope.renderer import RendererConfig, MatplotlibRenderer
from corrscope.layout import LayoutConfig

WIDTH = 640
HEIGHT = 360

ALL_ZEROS = np.array([0,0])

all_colors = pytest.mark.parametrize('bg_str,fg_str', [
    ('#000000', '#ffffff'),
    ('#ffffff', '#000000'),
    ('#0000aa', '#aaaa00'),
    ('#aaaa00', '#0000aa'),
])


@all_colors
def test_default_colors(bg_str, fg_str):
    """ Test the default background/foreground colors. """
    cfg = RendererConfig(
        WIDTH,
        HEIGHT,
        bg_color=bg_str,
        init_line_color=fg_str,
    )
    lcfg = LayoutConfig()
    nplots = 1

    r = MatplotlibRenderer(cfg, lcfg, nplots, None)
    verify(r, bg_str, fg_str)

    # Ensure default ChannelConfig(line_color=None) does not override line color
    chan = ChannelConfig(wav_path='')
    channels = [chan] * nplots
    r = MatplotlibRenderer(cfg, lcfg, nplots, channels)
    verify(r, bg_str, fg_str)


@all_colors
def test_line_colors(bg_str, fg_str):
    """ Test channel-specific line color overrides """
    cfg = RendererConfig(
        WIDTH,
        HEIGHT,
        bg_color=bg_str,
        init_line_color='#888888',
    )
    lcfg = LayoutConfig()
    nplots = 1

    chan = ChannelConfig(wav_path='', line_color=fg_str)
    channels = [chan] * nplots
    r = MatplotlibRenderer(cfg, lcfg, nplots, channels)
    verify(r, bg_str, fg_str)


def verify(r: MatplotlibRenderer, bg_str, fg_str):
    r.render_frame([ALL_ZEROS])
    frame_colors: np.ndarray = \
        np.frombuffer(r.get_frame(), dtype=np.uint8).reshape((-1, RGB_DEPTH))

    bg_u8 = [round(c*255) for c in to_rgb(bg_str)]
    fg_u8 = [round(c*255) for c in to_rgb(fg_str)]

    # Ensure background is correct
    assert (frame_colors[0] == bg_u8).all()

    # Ensure foreground is present
    assert np.prod(frame_colors == fg_u8, axis=-1).any(), \
        'incorrect foreground, it might be 136 = #888888'

    assert (np.amax(frame_colors, axis=0) == np.maximum(bg_u8, fg_u8)).all()
    assert (np.amin(frame_colors, axis=0) == np.minimum(bg_u8, fg_u8)).all()