diff --git a/aidge_core/unit_tests/test_show_graphview.py b/aidge_core/unit_tests/test_show_graphview.py index a14fc80bea101966a27c6050648a89a5a481340e..8cf2b0807a590ebd43409d573f6a0e0cd6c15ae7 100644 --- a/aidge_core/unit_tests/test_show_graphview.py +++ b/aidge_core/unit_tests/test_show_graphview.py @@ -3,6 +3,7 @@ import tempfile import unittest import builtins import aidge_core +import numpy as np from pathlib import Path from aidge_core.show_graphview import gview_to_json, str_aidge_graph_structure, str_aidge_seq_scheduling @@ -105,7 +106,26 @@ class test_show_gview(unittest.TestCase): if not hasattr(node_gview.get_operator(), 'get_micro_graph'): try: self.assertEqual(len(node_gview.get_operator().attr.dict()), len(node_json['attributes'])) - self.assertDictEqual(node_gview.get_operator().attr.dict(), node_json['attributes']) + + temp_node_dict = node_gview.get_operator().attr.dict() + + for key, value in node_gview.get_operator().attr.dict().items(): + + if isinstance(value, aidge_core.aidge_core.Tensor): + new_value = { + "dims": value.dims(), + "data_type": value.dtype(), + "tensor_data": np.array(value).tolist() + } + temp_node_dict.update({key : new_value}) + + elif not type(value).__name__ in dir(builtins): + temp_node_dict.update({key : str(value)}) + + else: + pass + + self.assertDictEqual(temp_node_dict, node_json['attributes']) except AttributeError: self.assertIsNone(node_gview.get_operator().attr) and self.assertFalse(node_json['attributes']) @@ -117,14 +137,24 @@ class test_show_gview(unittest.TestCase): for micro_node_gview in node_gview.get_operator().get_micro_graph().get_nodes(): for micro_node_json in node_json['attributes']['micro_graph']: if micro_node_gview.get_operator().type() == micro_node_json['optype']: - - for key, value in micro_node_gview.get_operator().attr.dict().items(): - if not type(value).__name__ in dir(builtins): - # Replace original value by its name (str) because value is of a type that could not be written to the JSON - # Cannot update this dict inplace : micro_node_gview.get_operator().attr.dict().update({key : value.name}) - temp_mnode_dict = micro_node_gview.get_operator().attr.dict() - temp_mnode_dict.update({key : value.name}) - self.assertDictEqual(temp_mnode_dict, micro_node_json['attributes']) + temp_mnode_dict = micro_node_gview.get_operator().attr.dict() # So the dict can be updated if needed + for key, value in micro_node_gview.get_operator().attr.dict().items(): + if isinstance(value, aidge_core.aidge_core.Tensor): + new_value = { + "dims": value.dims(), + "data_type": str(value.dtype()), + "tensor_data": np.array(value).tolist() + } + temp_mnode_dict.update({key : new_value}) + + elif not type(value).__name__ in dir(builtins): + # Use str(value) to stay consistent with how json.dumps(..., default=str) handles custom objects + temp_mnode_dict.update({key : str(value)}) + + else: + pass + + self.assertDictEqual(temp_mnode_dict, micro_node_json['attributes'])