pojagi-dsp/tests/channel/pipeline_test.py

import math
import pytest
from pojagi_dsp.channel import (
    ASignal,
    Constantly,
    Filter,
    FilterFunction,
    SignalFunction,
    IllegalStateException,
    Reduce,
)
from pojagi_dsp.channel.generator.sine import SineWave


@pytest.fixture
def srate():
    return 44100


@pytest.fixture
def freq():
    return 440


@pytest.fixture
def const(srate):
    return Constantly(42, srate=srate)


@pytest.fixture
def sine(srate, freq):
    return SineWave(freq, srate=srate)


@pytest.fixture
def sine_generator_factory(srate, freq):
    def sine():
        phase = 0.0
        inc = (2 * math.pi * freq) / srate
        while True:
            yield math.sin(phase)
            phase += inc

    return sine


def test_missing_srate_assigned_from_kwarg_reader(const: ASignal):
    pipeline = Filter(reader=const)
    assert pipeline.srate == const.srate


def test_missing_srate_assigned_from_ored_reader(const: ASignal):
    pipeline = const | Filter
    assert pipeline.srate == const.srate


def test_add_operator(const: Constantly):
    cursor = (const + 1).stream()
    assert next(cursor) == const.constant + 1


def test_mul_operator(const: Constantly):
    cursor = (const * 100).stream()
    assert next(cursor) == const.constant * 100


def test_filter_iterable(const: Constantly):
    pipeline = const | (Filter, Filter, Filter)
    assert pipeline.reader.reader.reader == const
    assert next(pipeline) == const.constant == next(const)


def test_filter_expression(const: Constantly):
    pipeline = const | (Filter | Filter | Filter)
    assert pipeline.reader.reader.reader == const
    assert next(pipeline) == const.constant == next(const)


def test_filter_nested_expression(const: Constantly):
    pipeline = const | (Filter | (Filter, Filter) | (Filter | Filter) | Filter)
    assert pipeline.reader.reader.reader.reader.reader.reader == const
    assert next(pipeline) == const.constant == next(const)


def test_reader(const: Constantly):
    filter = Filter()
    with pytest.raises(IllegalStateException, match=".reader. is None"):
        filter.reader
    assert (const | filter).reader == const


def test_pipeline_missing_reader(const: ASignal):
    pipeline = Filter | Filter | Filter
    with pytest.raises(IllegalStateException, match=".reader. is None"):
        next(pipeline)
    assert next(const | pipeline)


def test_filter_can_only_be_assigned_one_generator(const: Constantly):
    pipeline = const | Filter | Filter

    with pytest.raises(ValueError, match="already has a generator"):
        Constantly(0) | pipeline


def test_add_tuple(const: Constantly):
    pipeline = const + (100, 200, 300)
    assert isinstance(pipeline, Reduce)
    assert next(pipeline) == const.constant + (100 + 200 + 300)


def test_mul_tuple(const: Constantly):
    pipeline = const * (100, 200, 300)
    assert isinstance(pipeline, Reduce)
    assert next(pipeline) == const.constant * (100 * 200 * 300)


def test_radd(const: Constantly):
    pipeline = 1 + const
    assert next(pipeline) == const.constant + 1


def test_radd_tuple(const: Constantly):
    pipeline = (1, 2, 3) + const
    assert next(pipeline) == const.constant + (1 + 2 + 3)


def test_rmul(const: Constantly):
    pipeline = 2 * const
    assert next(pipeline) == const.constant * 2


def test_rmul_tuple(const: Constantly):
    pipeline = (2, 3, 4) * const
    assert next(pipeline) == const.constant * (2 * 3 * 4)


def test_samples(sine: SineWave, sine_generator_factory):
    pipeline = sine | Filter
    cursor = pipeline.samples()
    cursor2 = sine_generator_factory()
    for val in cursor:
        assert val == next(cursor2)


def test_stream(sine: SineWave, sine_generator_factory):
    pipeline = sine | Filter
    cursor = pipeline.stream()
    cursor2 = sine_generator_factory()
    for _ in range(5):
        assert next(cursor) == next(cursor2)


def test_cursor(sine: SineWave, sine_generator_factory):
    pipeline = sine | Filter
    cursor1 = pipeline.cursor
    cursor2 = pipeline.cursor
    cursor3 = sine_generator_factory()
    assert cursor1 == cursor2
    for v1 in cursor1:
        assert v1 == next(cursor3)

    with pytest.raises(StopIteration):
        next(cursor2)


def test_iter(sine: SineWave):
    count = 0
    for _ in sine:
        count += 1
    assert count == math.ceil(sine.wavelength)
    for _ in sine:
        count += 1
    assert count == math.ceil(sine.wavelength * 2)


def test_next(sine: SineWave):
    pipeline = sine | Filter
    for _ in range(5):
        val = next(pipeline)
        assert type(val) == float
        assert val != next(pipeline)


def test_function_filter(const: Constantly):
    pipeline = const | FilterFunction(lambda r: (x + 1 for x in r))
    assert next(pipeline) == const.constant + 1


def test_filter_function_literal(const: Constantly):
    pipeline = const | (lambda r: (x + 1 for x in r))
    assert next(pipeline) == const.constant + 1


def test_signal_function(sine: SineWave, sine_generator_factory):
    pipeline = SignalFunction(lambda _: sine.samples(), srate=srate) | Filter
    sine_generator = sine_generator_factory()
    for _ in range(5):
        assert next(sine_generator) == next(pipeline)


def test_sine_mod(sine: SineWave):
    seconds = 10
    stream = sine.stream()
    for _ in range(sine.srate * seconds):
        next(stream)
        assert sine.phase <= (2 * math.pi)

    for _ in sine:
        assert sine.phase <= (2 * math.pi)