Make Fit a line a normal unittest

python/paddle/v2/fluid/tests/book/test_fit_a_line.py

@@ -12,9 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 import paddle.v2 as paddle
 import paddle.v2.fluid as fluid
+import contextlib
+import unittest
+
+
+def main(use_cuda):
+    if use_cuda and not fluid.core.is_compiled_with_cuda():
+        return
 
     x = fluid.layers.data(name='x', shape=[13], dtype='float32')
 
@@ -35,7 +41,7 @@ train_reader = paddle.batch(
             paddle.dataset.uci_housing.train(), buf_size=500),
         batch_size=BATCH_SIZE)
 
-place = fluid.CPUPlace()
+    place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
     feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
     exe = fluid.Executor(place)
 
@@ -51,5 +57,29 @@ for pass_id in range(PASS_NUM):
                                       fetch_list=[avg_cost])
             print(avg_loss_value)
             if avg_loss_value[0] < 10.0:
-            exit(0)  # if avg cost less than 10.0, we think our code is good.
+                return
-exit(1)
+    raise AssertionError("Fit a line cost is too large, {0:2.2}".format(
+        avg_loss_value[0]))
+
+
+class TestFitALine(unittest.TestCase):
+    def test_cpu(self):
+        with self.program_scope_guard():
+            main(use_cuda=False)
+
+    def test_cuda(self):
+        with self.program_scope_guard():
+            main(use_cuda=True)
+
+    @contextlib.contextmanager
+    def program_scope_guard(self):
+        prog = fluid.Program()
+        startup_prog = fluid.Program()
+        scope = fluid.core.Scope()
+        with fluid.scope_guard(scope):
+            with fluid.program_guard(prog, startup_prog):
+                yield
+
+
+if __name__ == '__main__':
+    unittest.main()