refine test

Author: cangtianhuang

test/legacy_test/test_broadcast_tensors_op.py

@@ -16,7 +16,8 @@
 import unittest
 
 import numpy as np
-from op_test import OpTest, convert_float_to_uint16, paddle_static_guard
+from op_test import OpTest, convert_float_to_uint16
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle.base import core
@@ -260,7 +261,7 @@ def setUp(self):
     def test_api(self):
 
         def test_static():
-            with paddle_static_guard():
+            with static_guard():
                 with paddle.static.program_guard(paddle.static.Program()):
                     inputs = [
                         paddle.static.data(
@@ -273,8 +274,7 @@ def test_static():
                     paddle.broadcast_tensors(inputs)
 
         def test_dynamic():
-            paddle.disable_static()
-            try:
+            with dygraph_guard():
                 inputs = [
                     paddle.to_tensor(
                         np.random.random([4, 1, 4, 1]).astype(self.dtype)
@@ -294,8 +294,6 @@ def test_dynamic():
                     ),
                 ]
                 paddle.broadcast_tensors(inputs)
-            finally:
-                paddle.enable_static()
 
         test_static()
         test_dynamic()
@@ -314,7 +312,7 @@ def setUp(self):
 class TestRaiseBroadcastTensorsError(unittest.TestCase):
     def test_errors(self):
         def test_type():
-            with paddle_static_guard():
+            with static_guard():
                 with paddle.static.program_guard(paddle.static.Program()):
                     inputs = [
                         paddle.static.data(
@@ -327,7 +325,7 @@ def test_type():
                     paddle.broadcast_tensors(inputs)
 
         def test_dtype():
-            with paddle_static_guard():
+            with static_guard():
                 with paddle.static.program_guard(paddle.static.Program()):
                     inputs = [
                         paddle.static.data(
@@ -340,7 +338,7 @@ def test_dtype():
                     paddle.broadcast_tensors(inputs)
 
         def test_bcast_semantics():
-            with paddle_static_guard():
+            with static_guard():
                 with paddle.static.program_guard(paddle.static.Program()):
                     inputs = [
                         paddle.static.data(
@@ -353,7 +351,7 @@ def test_bcast_semantics():
                     paddle.broadcast_tensors(inputs)
 
         def test_bcast_semantics_complex64():
-            with paddle_static_guard():
+            with static_guard():
                 with paddle.static.program_guard(paddle.static.Program()):
                     inputs = [
                         paddle.static.data(
@@ -382,8 +380,7 @@ def test_bcast_semantics_complex64():
 class TestRaiseBroadcastTensorsErrorDyGraph(unittest.TestCase):
     def test_errors(self):
         def test_type():
-            paddle.disable_static()
-            try:
+            with dygraph_guard():
                 inputs = [
                     paddle.to_tensor(
                         np.ones(shape=[1, 1, 1, 1], dtype='float32', name="x4")
@@ -393,12 +390,9 @@ def test_type():
                     ),
                 ]
                 paddle.broadcast_tensors(inputs)
-            finally:
-                paddle.enable_static()
 
         def test_dtype():
-            paddle.disable_static()
-            try:
+            with dygraph_guard():
                 inputs = [
                     paddle.to_tensor(
                         np.ones(shape=[1, 1, 1, 1], dtype='int8', name="x6")
@@ -408,12 +402,9 @@ def test_dtype():
                     ),
                 ]
                 paddle.broadcast_tensors(inputs)
-            finally:
-                paddle.enable_static()
 
         def test_bcast_semantics():
-            paddle.disable_static()
-            try:
+            with dygraph_guard():
                 inputs = [
                     paddle.to_tensor(
                         np.ones(shape=[1, 3, 1, 1], dtype='float32', name="x9")
@@ -423,8 +414,6 @@ def test_bcast_semantics():
                     ),
                 ]
                 paddle.broadcast_tensors(inputs)
-            finally:
-                paddle.enable_static()
 
         self.assertRaises(TypeError, test_type)
         self.assertRaises(TypeError, test_dtype)
@@ -440,7 +429,7 @@ def set_dtypes(self):
         pass
 
     def test_single_static(self):
-        with paddle_static_guard():
+        with static_guard():
             with paddle.static.program_guard(paddle.static.Program()):
                 inputs = [
                     paddle.static.data(
@@ -451,17 +440,14 @@ def test_single_static(self):
                 self.assertEqual(len(outputs), 1)
 
     def test_single_dynamic(self):
-        paddle.disable_static()
-        try:
+        with dygraph_guard():
             inputs = [
                 paddle.to_tensor(
                     np.random.random([1, 4, 1, 4]).astype(self.dtype)
                 ),
             ]
             outputs = paddle.broadcast_tensors(inputs)
             self.assertEqual(len(outputs), 1)
-        finally:
-            paddle.enable_static()
 
 
 class TestBroadcastTensorsAPIZeroSize(unittest.TestCase):
@@ -476,7 +462,7 @@ def set_dtype(self):
         pass
 
     def test_zero_size_static(self):
-        with paddle_static_guard():
+        with static_guard():
             with paddle.static.program_guard(paddle.static.Program()):
                 inputs = [
                     paddle.static.data(
@@ -491,8 +477,7 @@ def test_zero_size_static(self):
                 self.assertEqual(outputs[1].shape, self.expected_shape)
 
     def test_zero_size_dynamic(self):
-        paddle.disable_static()
-        try:
+        with dygraph_guard():
             data1 = np.zeros(self.shape1, dtype=self.dtype)
             data2 = np.zeros(self.shape2, dtype=self.dtype)
 
@@ -503,8 +488,6 @@ def test_zero_size_dynamic(self):
             outputs = paddle.broadcast_tensors(inputs)
             self.assertEqual(outputs[0].shape, self.expected_shape)
             self.assertEqual(outputs[1].shape, self.expected_shape)
-        finally:
-            paddle.enable_static()
 
 
 class TestBroadcastTensorsAPIZeroSize_bool(TestBroadcastTensorsAPIZeroSize):