# Copyright 2025 XMOS LIMITED. # This Software is subject to the terms of the XMOS Public Licence: Version 1. from pathlib import Path import annotated_types from audio_dsp.design.parse_json import Graph, make_pipeline, insert_forks, DspJson, pipeline_to_dspjson from audio_dsp.models.stage import all_models from audio_dsp.stages import all_stages from typing import get_origin, get_args, Literal from types import UnionType def find_autoforks(graph): for node in graph.nodes: if 'AutoFork' in node.placement.name: return True assert False, "No AutoFork node found in the graph after insert_forks." def test_no_shared_edge(): json_str = """ { "name": "No shared edge", "fs": 44100, "inputs": [{ "name": "inputs", "channels": 1 }], "nodes": [ { "op_type": "VolumeControl", "config": {}, "placement": { "input": [["inputs", 0]], "name": "VolumeControl_1", "thread": 0 } } ], "outputs": [{ "name": "outputs", "input": [["VolumeControl_1", 0]] }] } """ graph = Graph.model_validate_json(json_str) dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph) a = make_pipeline(dsp_json) a.draw(Path("test_no_shared_edge")) new_graph = insert_forks(graph) print(f"Before insert_forks: {graph.model_dump_json()}") print(f"After insert_forks: {new_graph.model_dump_json()}") dsp_json = DspJson(ir_version=1, producer_name="pipeline_to_dspjson", producer_version="1.0", graph=new_graph) new_json = pipeline_to_dspjson(a) assert dsp_json.graph == new_json.graph, "Pipeline JSON does not match original" for node in graph.nodes: if 'AutoFork' in node.placement.name: assert False, "AutoFork node found in the graph after insert_forks, but not needed." def test_shared_edge_from_graph_input(): json_str = """ { "name": "Shared edge from graph input", "fs": 44100, "inputs": [{ "name": "inputs", "channels": 1 }], "nodes": [ { "op_type": "VolumeControl", "config": {}, "placement": { "input": [["inputs", 0]], "name": "VolumeControl_A", "thread": 0 } }, { "op_type": "Mixer", "config": {}, "placement": { "input": [["inputs", 0]], "name": "Mixer_B", "thread": 0 } } ], "outputs": [{ "name": "outputs", "input": [["VolumeControl_A", 0]] }] } """ graph = Graph.model_validate_json(json_str) dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph) a = make_pipeline(dsp_json) a.draw(Path("test_shared_edge_from_graph_input")) new_graph = insert_forks(graph) print(f"Before insert_forks: {graph.model_dump_json()}") print(f"After insert_forks: {new_graph.model_dump_json()}") find_autoforks(new_graph) def test_shared_edge_from_producer_node(): json_str = """ { "name": "Shared edge from producer node", "fs": 44100, "inputs": [{ "name": "inputs", "channels": 1 }], "nodes": [ { "op_type": "VolumeControl", "config": {}, "placement": { "input": [["inputs", 0]], "name": "Producer", "thread": 0 } }, { "op_type": "Mixer", "config": {}, "placement": { "input": [["Producer", 0]], "name": "Consumer_1", "thread": 0 } }, { "op_type": "Mixer", "config": {}, "placement": { "input": [["Producer", 0]], "name": "Consumer_2", "thread": 0 } } ], "outputs": [{ "name": "outputs", "input": [["Consumer_1", 0]] }] } """ graph = Graph.model_validate_json(json_str) dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph) a = make_pipeline(dsp_json) a.draw(Path("test_shared_edge_from_producer_node")) new_graph = insert_forks(graph) print(f"Before insert_forks: {graph.model_dump_json()}") print(f"After insert_forks: {new_graph.model_dump_json()}") find_autoforks(new_graph) def test_shared_edge_with_graph_output(): json_str = """ { "name": "Shared edge with graph output", "fs": 44100, "inputs": [{ "name": "inputs", "channels": 1 }], "nodes": [ { "op_type": "VolumeControl", "config": {}, "placement": { "input": [["inputs", 0]], "name": "Producer", "thread": 0 } }, { "op_type": "Mixer", "config": {}, "placement": { "input": [["Producer", 0]], "name": "Consumer", "thread": 0 } } ], "outputs": [{ "name": "outputs", "input": [["Producer", 0]] }] } """ graph = Graph.model_validate_json(json_str) dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph) a = make_pipeline(dsp_json) a.draw(Path("test_shared_edge_with_graph_output")) new_graph = insert_forks(graph) print(f"Before insert_forks: {graph.model_dump_json()}") print(f"After insert_forks: {new_graph.model_dump_json()}") find_autoforks(new_graph) def test_again(): json_str = """ { "name": "Stereo Mixer with Volume", "fs": 48000, "inputs": [{ "name": "stereo_in", "channels": 2 }], "nodes": [ {"op_type": "Mixer", "placement": {"input": [["stereo_in", 0], ["stereo_in", 1]], "name": "Mixer", "thread": 0}}, {"op_type": "VolumeControl", "placement": {"input": [["Mixer", 0]], "name": "Volume", "thread": 0}} ], "outputs": [{ "name": "stereo_out", "input": [["Volume", 0], ["Volume", 0]] }] } """ graph = Graph.model_validate_json(json_str) dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph) a = make_pipeline(dsp_json) a.draw(Path("test_again")) new_graph = insert_forks(graph) print(f"Before insert_forks: {graph.model_dump_json()}") print(f"After insert_forks: {new_graph.model_dump_json()}") find_autoforks(new_graph) def test_multiple_inputs_outputs_non_shared(): """ Test a graph with two inputs (one mono and one stereo) and two outputs (one mono and one stereo) where no edge is shared between consumers. """ json_str = """ { "name": "Multiple Inputs and Outputs Non-Shared Test", "fs": 48000, "inputs": [ {"name": "mono_in", "channels": 1}, {"name": "stereo_in", "channels": 2} ], "nodes": [ { "op_type": "VolumeControl", "config": {}, "placement": { "input": [["mono_in", 0]], "name": "VolumeControl_A", "thread": 0 } }, { "op_type": "Mixer", "config": {}, "placement": { "input": [["stereo_in", 0], ["stereo_in", 1]], "name": "Mixer_B", "thread": 0 } } ], "outputs": [ {"name": "mono_out", "input": [["VolumeControl_A", 0]]}, {"name": "stereo_out", "input": [["Mixer_B", 0], ["VolumeControl_A", 0]]} ] } """ graph = Graph.model_validate_json(json_str) dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph) a = make_pipeline(dsp_json) a.draw(Path("test_multiple_inputs_outputs_non_shared")) new_graph = insert_forks(graph) print("Test: Multiple Inputs and Outputs Non-Shared Test") print(f"Before insert_forks: {graph.model_dump_json(indent=2)}") print(f"After insert_forks: {new_graph.model_dump_json(indent=2)}") find_autoforks(new_graph) def test_multiple_inputs_outputs_shared(): """ Test a graph with two inputs and two outputs where a node produces an output that is shared by both graph outputs. A Fork should be inserted. """ json_str = """ { "name": "Multiple Inputs and Outputs Shared Test", "fs": 48000, "inputs": [ {"name": "input1", "channels": 1}, {"name": "input2", "channels": 1} ], "nodes": [ { "op_type": "Mixer", "config": {}, "placement": { "input": [["input1", 0], ["input2", 0]], "name": "Mixer", "thread": 0 } } ], "outputs": [ {"name": "output1", "input": [["Mixer", 0]]}, {"name": "output2", "input": [["Mixer", 0]]} ] } """ graph = Graph.model_validate_json(json_str) dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph) a = make_pipeline(dsp_json) a.draw(Path("test_multiple_inputs_outputs_shared")) new_graph = insert_forks(graph) print("Test: Multiple Inputs and Outputs Shared Test") print(f"Before insert_forks: {graph.model_dump_json(indent=2)}") print(f"After insert_forks: {new_graph.model_dump_json(indent=2)}") find_autoforks(new_graph) def test_all_stages_models(): all_m = all_models() all_s = all_stages() failed = False for s in all_s: if s.startswith("_") or s in ["DSPThreadStage", "PipelineStage"]: continue try: assert s in all_m, f"Stage {s} not found in all models" assert s == all_m[s].model_fields["op_type"].default, f"Stage {s} op_type mismatch" except AssertionError as e: print(e) failed = True continue if "biquad" in s.lower() or "parametric" in s.lower(): continue if "parameters" not in all_m[s].model_fields: continue try: if type(all_s[s].set_parameters.__annotations__["parameters"]) is UnionType: set_params_input = get_args(all_s[s].set_parameters.__annotations__["parameters"]) model_params_type = all_m[s].model_fields["parameters"].default_factory assert model_params_type in set_params_input, f"Stage {s} set_parameters input type mismatch" else: set_params_input = all_s[s].set_parameters.__annotations__["parameters"].__name__ model_params_type = all_m[s].model_fields["parameters"].default_factory.__name__ assert set_params_input == model_params_type, f"Stage {s} set_parameters input type mismatch" except AssertionError as e: print(e) failed = True for field, value in type(all_m[s].model_fields["parameters"].default_factory()).model_fields.items(): try: if get_origin(value.annotation) is list: item_type = get_args(value.annotation)[0] meta = get_args(item_type)[1].metadata elif get_origin(value.annotation) is Literal: continue # Skip Literal types else: meta = value.metadata min = [g.ge for g in meta if isinstance(g, annotated_types.Ge)] or [ g.gt for g in meta if isinstance(g, annotated_types.Gt)] max = [g.le for g in meta if isinstance(g, annotated_types.Le)] or [ g.lt for g in meta if isinstance(g, annotated_types.Lt)] assert min, f"Minimum not defined for field {field} in {s}" if field not in ["position", "delay"]: assert max, f"Maximum not defined for field {field} in {s}" assert min[0] < max[0], f"Range not correct for field {field} in {s}" except AssertionError as e: print(e) failed = True continue assert not failed, "Some stages failed the test." if __name__ == "__main__": test_shared_edge_from_graph_input()