Add df_service

suaaa7 · suaaa7 · commit f6193ddd177f · 2020-10-13T00:02:14.000+09:00
diff --git a/python_ci/requirements.txt b/python_ci/requirements.txt
@@ -1 +1,2 @@
 requests==2.24.0
+pandas==1.1.3
diff --git a/python_ci/src/df_service.py b/python_ci/src/df_service.py
@@ -0,0 +1,36 @@
+from typing import Optional
+
+import pandas as pd
+
+
+class DFService:
+    def swap_and_remove_path(
+        self,
+        df: pd.DataFrame,
+        base_dir: Optional[str] = None
+    ) -> pd.DataFrame:
+        service = ImageDownloadService()
+
+        for i, image_hashes in enumerate(df.image_hash):
+            image_paths = []
+            for image_hash in image_hashes:
+                maybe_local_path = service.get_local_path(
+                    image_hash=image_hash,
+                    base_dir=base_dir
+                )
+                if maybe_local_path is None:
+                    print("Skipped this data.")
+                else:
+                    image_paths.append(maybe_local_path)
+            df.at[i, 'image_path'] = image_paths
+
+        return df[df['image_path'].map(len) == df['num']]
+
+
+class ImageDownloadService:
+    def get_local_path(
+        self,
+        image_hash: str,
+        base_dir: Optional[str]
+    ) -> Optional[str]:
+        return None
diff --git a/python_ci/src/tests/test_df_service.py b/python_ci/src/tests/test_df_service.py
@@ -0,0 +1,101 @@
+from unittest import TestCase, main
+from unittest.mock import patch
+
+import pandas as pd
+from pandas.testing import assert_frame_equal
+
+from src.df_service import DFService
+
+
+class TestDFService(TestCase):
+    def setUp(self) -> None:
+        self.patcher = patch('src.df_service.ImageDownloadService.get_local_path')
+        self.m_get_local_path = self.patcher.start()
+
+        self.service = DFService()
+        self.test_df = pd.DataFrame({
+            'num': [1, 2],
+            'image_hash': [
+                ['aaaaa'],
+                ['aaaaa', 'aaaaa']
+            ],
+            'image_path': [
+                ['s3://test_image.jpg'],
+                ['s3://test_image.jpg', 's3://test_image.jpg']
+            ]
+        })
+        self.test_base_dir = '/home/test-user'
+
+    def tearDown(self) -> None:
+        self.patcher.stop()
+
+    def test_return_empty_df(self) -> None:
+        self.m_get_local_path.return_value = None
+
+        actual_df = self.service.swap_and_remove_path(
+            self.test_df,
+            self.test_base_dir
+        )
+
+        self.assertTrue(actual_df.empty)
+        self.assertEqual(self.m_get_local_path.call_count, 3)
+        self.m_get_local_path.assert_called_with(
+            image_hash='aaaaa',
+            base_dir=self.test_base_dir
+        )
+
+    def test_return_empty_df_when_base_dir_is_none(self) -> None:
+        self.m_get_local_path.return_value = None
+
+        actual_df = self.service.swap_and_remove_path(self.test_df)
+
+        self.assertTrue(actual_df.empty)
+        self.assertEqual(self.m_get_local_path.call_count, 3)
+        self.m_get_local_path.assert_called_with(
+            image_hash='aaaaa',
+            base_dir=None
+        )
+
+    def test_return_expected_df(self) -> None:
+        self.m_get_local_path.return_value = '/opt/program/test_image.jpg'
+
+        expected_df = pd.DataFrame({
+            'num': [1, 2],
+            'image_hash': [
+                ['aaaaa'],
+                ['aaaaa', 'aaaaa']
+            ],
+            'image_path': [
+                ['/opt/program/test_image.jpg'],
+                ['/opt/program/test_image.jpg', '/opt/program/test_image.jpg']
+            ]
+        })
+        actual_df = self.service.swap_and_remove_path(self.test_df)
+
+        assert_frame_equal(actual_df, expected_df)
+        self.assertEqual(self.m_get_local_path.call_count, 3)
+
+    def test_return_expected_df_when_get_local_path_return_none(self) -> None:
+        self.m_get_local_path.side_effect = [
+            '/opt/program/test_image.jpg',
+            None,
+            '/opt/program/test_image.jpg'
+        ]
+
+        expected_df = pd.DataFrame({
+            'num': [1],
+            'image_hash': [
+                ['aaaaa']
+            ],
+            'image_path': [
+                ['/opt/program/test_image.jpg']
+            ]
+        })
+        actual_df = self.service.swap_and_remove_path(self.test_df)
+
+        assert_frame_equal(actual_df, expected_df)
+        self.assertEqual(self.m_get_local_path.call_count, 3)
+
+
+if __name__ == '__main__':
+    main()
