fix(tests): refactoring video tests

jina-ai · Apr 17, 2020 · b3b8e93 · b3b8e93
1 parent ead2d48
commit b3b8e93
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 84 deletions.
diff --git a/tests/executors/encoders/image/__init__.py b/tests/executors/encoders/image/__init__.py
@@ -20,17 +20,34 @@ def target_output_dim(self):
     def target_output_dim(self, output_dim):
         self._target_output_dim = output_dim
 
+    @property
+    def input_dim(self):
+        return self._input_dim
+
+    @input_dim.setter
+    def input_dim(self, input_dim):
+        self._input_dim = input_dim
+
+    def get_encoder(self):
+        encoder = self._get_encoder()
+        encoder.workspace = self.workspace
+        self.add_tmpfile(encoder.workspace)
+        return encoder
+
+    def _get_encoder(self):
+        raise NotImplementedError
+
     @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
     def test_encoding_results(self):
         encoder = self.get_encoder()
-        test_data = np.random.rand(2, 3, 224, 224)
+        test_data = np.random.rand(2, 3, self.input_dim, self.input_dim)
         encoded_data = encoder.encode(test_data)
-        self.assertEqual(encoded_data.shape, (2, self._target_output_dim))
+        self.assertEqual(encoded_data.shape, (2, self.target_output_dim))
 
     @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
     def test_save_and_load(self):
         encoder = self.get_encoder()
-        test_data = np.random.rand(2, 3, 224, 224)
+        test_data = np.random.rand(2, 3, self.input_dim, self.input_dim)
         encoded_data_control = encoder.encode(test_data)
         encoder.touch()
         encoder.save()
@@ -47,12 +64,3 @@ def test_save_and_load_config(self):
         self.assertTrue(os.path.exists(encoder.config_abspath))
         encoder_loaded = BaseExecutor.load_config(encoder.config_abspath)
         self.assertEqual(encoder_loaded.pool_strategy, encoder.pool_strategy)
-
-    def get_encoder(self):
-        encoder = self._get_encoder()
-        encoder.workspace = self.workspace
-        self.add_tmpfile(encoder.workspace)
-        return encoder
-
-    def _get_encoder(self):
-        raise NotImplementedError
diff --git a/tests/executors/encoders/video/__init__.py b/tests/executors/encoders/video/__init__.py
@@ -0,0 +1,58 @@
+import os
+import unittest
+
+import numpy as np
+
+from jina.executors import BaseExecutor
+from tests import JinaTestCase
+
+
+class VideoTestCase(JinaTestCase):
+    @property
+    def workspace(self):
+        return os.path.join(os.environ['TEST_WORKDIR'], 'test_image')
+
+    @property
+    def target_output_dim(self):
+        return self._target_output_dim
+
+    @target_output_dim.setter
+    def target_output_dim(self, output_dim):
+        self._target_output_dim = output_dim
+
+    def get_encoder(self):
+        encoder = self._get_encoder()
+        encoder.workspace = self.workspace
+        self.add_tmpfile(encoder.workspace)
+        return encoder
+
+    def _get_encoder(self):
+        raise NotImplementedError
+
+    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
+    def test_encoding_results(self):
+        encoder = self.get_encoder()
+        test_data = np.random.rand(2, 3, 3, 224, 224)
+        encoded_data = encoder.encode(test_data)
+        self.assertEqual(encoded_data.shape, (2, self._target_output_dim))
+
+    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
+    def test_save_and_load(self):
+        encoder = self.get_encoder()
+        test_data = np.random.rand(2, 3, 3, 224, 224)
+        encoded_data_control = encoder.encode(test_data)
+        encoder.touch()
+        encoder.save()
+        self.assertTrue(os.path.exists(encoder.save_abspath))
+        encoder_loaded = BaseExecutor.load(encoder.save_abspath)
+        encoded_data_test = encoder_loaded.encode(test_data)
+        self.assertEqual(encoder_loaded.model_name, encoder.model_name)
+        np.testing.assert_array_equal(encoded_data_control, encoded_data_test)
+
+    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
+    def test_save_and_load_config(self):
+        encoder = self.get_encoder()
+        encoder.save_config()
+        self.assertTrue(os.path.exists(encoder.config_abspath))
+        encoder_loaded = BaseExecutor.load_config(encoder.config_abspath)
+        self.assertEqual(encoder_loaded.model_name, encoder.model_name)
diff --git a/tests/executors/encoders/video/test_paddlehub.py b/tests/executors/encoders/video/test_paddlehub.py
@@ -1,44 +1,14 @@
-import os
 import unittest
 
-import numpy as np
-
-from jina.executors import BaseExecutor
 from jina.executors.encoders.video.paddlehub import VideoPaddlehubEncoder
-from tests import JinaTestCase
-
-
-class MyTestCase(JinaTestCase):
-    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
-    def test_encoding_results(self):
-        encoder = VideoPaddlehubEncoder()
-        test_data = np.random.rand(2, 3, 3, 224, 224)
-        encoded_data = encoder.encode(test_data)
-        self.assertEqual(encoded_data.shape, (2, 2048))
+from . import VideoTestCase
 
-    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
-    def test_save_and_load(self):
-        encoder = VideoPaddlehubEncoder()
-        test_data = np.random.rand(2, 3, 3, 224, 224)
-        encoded_data_control = encoder.encode(test_data)
-        encoder.touch()
-        encoder.save()
-        self.assertTrue(os.path.exists(encoder.save_abspath))
-        encoder_loaded = BaseExecutor.load(encoder.save_abspath)
-        encoded_data_test = encoder_loaded.encode(test_data)
-        self.assertEqual(encoder_loaded.model_name, encoder.model_name)
-        np.testing.assert_array_equal(encoded_data_control, encoded_data_test)
-        self.add_tmpfile(
-            encoder.config_abspath, encoder.save_abspath, encoder_loaded.config_abspath, encoder_loaded.save_abspath)
 
-    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
-    def test_save_and_load_config(self):
-        encoder = VideoPaddlehubEncoder()
-        encoder.save_config()
-        self.assertTrue(os.path.exists(encoder.config_abspath))
-        encoder_loaded = BaseExecutor.load_config(encoder.config_abspath)
-        self.assertEqual(encoder_loaded.model_name, encoder.model_name)
-        self.add_tmpfile(encoder_loaded.config_abspath, encoder_loaded.save_abspath)
+class MyTestCase(VideoTestCase):
+    def _get_encoder(self):
+        self.target_output_dim = 2048
+        self.input_dim = 224
+        return VideoPaddlehubEncoder()
 
 
 if __name__ == '__main__':

diff --git a/tests/executors/encoders/video/test_torchvision.py b/tests/executors/encoders/video/test_torchvision.py
@@ -1,44 +1,14 @@
 import unittest
 
-import os
-import numpy as np
-
-from jina.executors import BaseExecutor
 from jina.executors.encoders.video.torchvision import VideoTorchEncoder
-from tests import JinaTestCase
-
-
-class MyTestCase(JinaTestCase):
-    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
-    def test_encoding_results(self):
-        encoder = VideoTorchEncoder()
-        test_data = np.random.rand(2, 3, 3, 112, 112)
-        encoded_data = encoder.encode(test_data)
-        self.assertEqual(encoded_data.shape, (2, 512))
+from . import VideoTestCase
 
-    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
-    def test_save_and_load(self):
-        encoder = VideoTorchEncoder()
-        test_data = np.random.rand(2, 3, 3, 112, 112)
-        encoded_data_control = encoder.encode(test_data)
-        encoder.touch()
-        encoder.save()
-        self.assertTrue(os.path.exists(encoder.save_abspath))
-        encoder_loaded = BaseExecutor.load(encoder.save_abspath)
-        encoded_data_test = encoder_loaded.encode(test_data)
-        self.assertEqual(encoder_loaded.model_name, encoder.model_name)
-        np.testing.assert_array_equal(encoded_data_control, encoded_data_test)
-        self.add_tmpfile(
-            encoder.config_abspath, encoder.save_abspath, encoder_loaded.config_abspath, encoder_loaded.save_abspath)
 
-    @unittest.skipUnless('JINA_TEST_PRETRAINED' in os.environ, 'skip the pretrained test if not set')
-    def test_save_and_load_config(self):
-        encoder = VideoTorchEncoder()
-        encoder.save_config()
-        self.assertTrue(os.path.exists(encoder.config_abspath))
-        encoder_loaded = BaseExecutor.load_config(encoder.config_abspath)
-        self.assertEqual(encoder_loaded.model_name, encoder.model_name)
-        self.add_tmpfile(encoder_loaded.config_abspath, encoder_loaded.save_abspath)
+class MyTestCase(VideoTestCase):
+    def _get_encoder(self):
+        self.target_output_dim = 512
+        self.input_dim = 112
+        return VideoTorchEncoder()
 
 
 if __name__ == '__main__':