compiler/visq-unittest/test/testDotBuilder.py

   1 # Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #    http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 """Test visqlib.DotBuilder module"""
  15
  16 import unittest
  17 import pydot
  18
  19 from visqlib.DotBuilder import DotBuilder
  20 from test.Resources import fp32_model_dir
  21
  22
  23 class VisqDotBuilderTest(unittest.TestCase):
  24     def test_dot_builder_wrong_input_file(self):
  25         self.assertRaises(FileNotFoundError, DotBuilder, "wrong", "wrong", "wrong",
  26                           "wrong")
  27
  28     def test_dot_builder(self):
  29         test_colors = [{"b": 0, "e": 0.5, "c": "green"}, {"b": 0.5, "e": 1, "c": "red"}]
  30         test_qerror_map = dict()
  31         test_qerror_map["ofm"] = 0.1
  32         builder = DotBuilder(fp32_model_dir + "/Add_000.circle", "Add_000.dot", "MPEIR",
  33                              test_colors)
  34         builder.save(test_qerror_map)
  35
  36         graph = pydot.graph_from_dot_file("Add_000.dot")[0]
  37         # Why 1? 0 is output
  38         ofm_node = graph.get_node("\"ofm\"")[1]
  39         self.assertEqual("green", ofm_node.get_fillcolor())
  40
  41
  42 if __name__ == "__main__":
  43     unittest.main()