improve test

tests/links_tests/model_tests/test_sequential_feature_extractor.py

@@ -1,14 +1,12 @@
-import unittest
-
 import numpy as np
+import unittest
 
 import chainer
 from chainer.cuda import to_cpu
+from chainer.function import Function
 from chainer import testing
 from chainer.testing import attr
 
-from chainer.function import Function
-
 from chainercv.links import SequentialFeatureExtractor
 from chainercv.utils.testing import ConstantStubLink
 
@@ -19,7 +17,13 @@ def forward(self, inputs):
         return inputs[0] * 2,
 
 
-class TestSequentialFeatureExtractorOrderedDictFunctions(unittest.TestCase):
+@testing.parameterize(
+    {'layer_names': None},
+    {'layer_names': 'f2'},
+    {'layer_names': ('f2',)},
+    {'layer_names': ('l2', 'l1', 'f2')},
+)
+class TestSequentialFeatureExtractor(unittest.TestCase):
 
     def setUp(self):
         self.l1 = ConstantStubLink(np.random.uniform(size=(1, 3, 24, 24)))
@@ -33,53 +37,48 @@ def setUp(self):
             self.link.f1 = self.f1
             self.link.f2 = self.f2
             self.link.l2 = self.l2
-        self.link.layer_names = ['l1', 'f1', 'f2']
+
+        if self.layer_names:
+            self.link.layer_names = self.layer_names
+
         self.x = np.random.uniform(size=(1, 3, 24, 24))
 
-    def check_call_output(self):
+    def check_call(self):
         x = self.link.xp.asarray(self.x)
-        out = self.link(x)
-
-        self.assertEqual(len(out), 3)
-        self.assertIsInstance(out[0], chainer.Variable)
-        self.assertIsInstance(out[1], chainer.Variable)
-        self.assertIsInstance(out[2], chainer.Variable)
-        self.assertIsInstance(out[0].data, self.link.xp.ndarray)
-        self.assertIsInstance(out[1].data, self.link.xp.ndarray)
-        self.assertIsInstance(out[2].data, self.link.xp.ndarray)
 
-        out_data = [to_cpu(var.data) for var in out]
-        np.testing.assert_equal(out_data[0], to_cpu(self.l1(x).data))
-        np.testing.assert_equal(out_data[1], to_cpu(self.f1(self.l1(x)).data))
-        np.testing.assert_equal(
-            out_data[2], to_cpu(self.f2(self.f1(self.l1(x))).data))
+        expects = dict()
+        expects['l1'] = self.l1(x)
+        expects['f1'] = self.f1(expects['l1'])
+        expects['f2'] = self.f2(expects['f1'])
+        expects['l2'] = self.l2(expects['f2'])
 
-    def test_call_output_cpu(self):
-        self.check_call_output()
+        outs = self.link(x)
 
-    @attr.gpu
-    def test_call_output_gpu(self):
-        self.link.to_gpu()
-        self.check_call_output()
+        if isinstance(self.layer_names, tuple):
+            layer_names = self.layer_names
+        else:
+            if self.layer_names is None:
+                layer_names = ('l2',)
+            else:
+                layer_names = (self.layer_names,)
+            outs = (outs,)
 
-    def check_call_dynamic_layer_names(self):
-        x = self.link.xp.asarray(self.x)
-        self.link.layer_names = ['l2']
-        out, = self.link(x)
+        self.assertEqual(len(outs), len(layer_names))
 
-        self.assertIsInstance(out, chainer.Variable)
-        self.assertIsInstance(out.data, self.link.xp.ndarray)
+        for out, layer_name in zip(outs, layer_names):
+            self.assertIsInstance(out, chainer.Variable)
 
-        out_data = out.data
-        np.testing.assert_equal(
-            out_data, to_cpu(self.l2(self.f2(self.f1(self.l1(x)))).data))
+            out = to_cpu(out.data)
+            self.assertIsInstance(out, self.link.xp.ndarray)
+            np.testing.assert_equal(out, to_cpu(expects[layer_name].data))
 
-    def test_call_dynamic_layer_names_cpu(self):
-        self.check_call_dynamic_layer_names()
+    def test_call_cpu(self):
+        self.check_call()
 
     @attr.gpu
-    def test_call_dynamic_layer_names_gpu(self):
-        self.check_call_dynamic_layer_names()
+    def test_call_gpu(self):
+        self.link.to_gpu()
+        self.check_call()
 
 
 testing.run_module(__name__, __file__)