make all test class derived from Tf2OnnxBackendTestBase

Author: nbcsm

tests/backend_test_base.py

@@ -26,20 +26,16 @@
 class Tf2OnnxBackendTestBase(unittest.TestCase):
     def setUp(self):
         self.config = get_test_config()
-        self.maxDiff = None
         tf.reset_default_graph()
         # reset name generation on every test
         utils.INTERNAL_NAME = 1
         np.random.seed(1)  # Make it reproducible.
 
         self.log = logging.getLogger("tf2onnx.unitest." + str(type(self)))
-        if self.config.is_debug_mode:
-            self.log.setLevel(logging.DEBUG)
-        else:
+        if not self.config.is_debug_mode:
             # suppress log info of tensorflow so that result of test can be seen much easier
             os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
             tf.logging.set_verbosity(tf.logging.WARN)
-            self.log.setLevel(logging.INFO)
 
     def tearDown(self):
         if not self.config.is_debug_mode:

tests/test_graph.py

@@ -7,8 +7,6 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import os
-import unittest
 from collections import namedtuple
 
 import numpy as np
@@ -17,13 +15,13 @@
 import tensorflow as tf
 from onnx import helper
 
-import tf2onnx
 from tf2onnx import constants, utils
 from tf2onnx.graph import GraphUtil
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.tfonnx import process_tf_graph
 from tf2onnx.handler import tf_op
 
+from backend_test_base import Tf2OnnxBackendTestBase
 from common import get_test_config, unittest_main, check_tf_min_version, check_tf_max_version
 
 
@@ -91,21 +89,12 @@ def onnx_pretty(g, args=None):
     return helper.printable_graph(model_proto.graph)
 
 
-class Tf2OnnxGraphTests(unittest.TestCase):
+class Tf2OnnxGraphTests(Tf2OnnxBackendTestBase):
     """Test cases."""
     maxDiff = None
 
     def setUp(self):
-        """Setup test."""
-        # reset name generation on every test
-        # suppress log info of tensorflow so that result of test can be seen much easier
-        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
-        tf.logging.set_verbosity(tf.logging.WARN)
-
-        self.config = get_test_config()
-
-        tf2onnx.utils.INTERNAL_NAME = 1
-        tf.reset_default_graph()
+        super().setUp()
         arg = namedtuple("Arg", "input inputs outputs verbose continue_on_error")
         self._args0 = arg(input="test", inputs=[], outputs=["output:0"],
                           verbose=False, continue_on_error=False)

tests/test_internals.py

@@ -7,8 +7,6 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import os
-import unittest
 from collections import namedtuple
 
 import graphviz as gv
@@ -20,6 +18,8 @@
 from tf2onnx import utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.graph import GraphUtil
+
+from backend_test_base import Tf2OnnxBackendTestBase
 from common import unittest_main
 
 
@@ -49,14 +49,9 @@ def onnx_pretty(g, args=None):
     return helper.printable_graph(graph_proto.graph)
 
 
-class Tf2OnnxInternalTests(unittest.TestCase):
+class Tf2OnnxInternalTests(Tf2OnnxBackendTestBase):
     def setUp(self):
-        """Setup test."""
-        # suppress log info of tensorflow so that result of test can be seen much easier
-        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
-        tf.logging.set_verbosity(tf.logging.WARN)
-
-        utils.INTERNAL_NAME = 1
+        super().setUp()
         arg = namedtuple("Arg", "input inputs outputs verbose")
         self._args0 = arg(input="test", inputs=[], outputs=["output:0"], verbose=False)
         self._args1 = arg(input="test", inputs=["input:0"], outputs=["output:0"], verbose=False)