diff --git a/tests/common/debugging/test_drawing.py b/tests/common/debugging/test_drawing.py new file mode 100644 index 000000000..1572996b9 --- /dev/null +++ b/tests/common/debugging/test_drawing.py @@ -0,0 +1,28 @@ +"""Test file for drawing""" + +import tempfile +from pathlib import Path + +from hdk.common.data_types.integers import Integer +from hdk.common.debugging import draw_graph +from hdk.common.values import EncryptedValue +from hdk.hnumpy.compile import compile_numpy_function_into_op_graph + + +def test_draw_graph_with_saving(): + """Tests drawing and saving a graph""" + + def function(x): + return x + 42 + + op_graph = compile_numpy_function_into_op_graph( + function, + {"x": EncryptedValue(Integer(7, True))}, + iter([(-2,), (-1,), (0,), (1,), (2,)]), + ) + + with tempfile.TemporaryDirectory() as tmp: + output_directory = Path(tmp) + output_file = output_directory.joinpath("test.png") + draw_graph(op_graph, save_to=output_file) + assert output_file.exists()