Finish tests/flow/extract

Author: Jimmy Li

tests/flow/extract/test_extract_html_flow.py

@@ -0,0 +1,51 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from uniflow.flow.extract.extract_html_flow import ExtractHTMLFlow
+from uniflow.node import Node
+from uniflow.op.extract.split.constants import PARAGRAPH_SPLITTER
+
+
+class TestExtractHTMLFlow(unittest.TestCase):
+    @patch("uniflow.flow.extract.extract_html_flow.ExtractHTMLOp")
+    @patch("uniflow.flow.extract.extract_html_flow.ProcessHTMLOp")
+    @patch("uniflow.flow.extract.extract_html_flow.SplitterOpsFactory")
+    def setUp(
+        self, mock_splitter_ops_factory, mock_process_html_op, mock_extract_html_op
+    ):
+        self.mock_splitter_ops_factory = mock_splitter_ops_factory
+        self.mock_process_html_op = mock_process_html_op
+        self.mock_extract_html_op = mock_extract_html_op
+        self.extract_html_flow = ExtractHTMLFlow()
+
+    def test_init(self):
+        self.mock_extract_html_op.assert_called_once_with(name="extract_html_op")
+        self.mock_process_html_op.assert_called_once_with(name="process_html_op")
+        self.mock_splitter_ops_factory.get.assert_called_once_with(PARAGRAPH_SPLITTER)
+
+    def test_run(self):
+        # arrange
+        nodes = [
+            Node(name="node1", value_dict={"filename": "filepath"}),
+            Node(name="node2", value_dict={"filename": "filepath"}),
+        ]
+
+        self.mock_splitter_ops_factory.get.return_value.return_value = MagicMock()
+        self.mock_process_html_op.return_value.return_value = MagicMock()
+        self.mock_extract_html_op.return_value.return_value = MagicMock()
+        result = self.extract_html_flow.run(nodes)
+
+        self.mock_extract_html_op.return_value.assert_called_once_with(nodes)
+        self.mock_process_html_op.return_value.assert_called_once_with(
+            self.mock_extract_html_op.return_value.return_value
+        )
+        self.mock_splitter_ops_factory.get.return_value.assert_called_once_with(
+            self.mock_process_html_op.return_value.return_value
+        )
+        self.assertEqual(
+            result, self.mock_splitter_ops_factory.get.return_value.return_value
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/flow/extract/test_extract_image_flow.py

@@ -0,0 +1,63 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from uniflow.flow.extract.extract_image_flow import ExtractImageFlow
+from uniflow.node import Node
+from uniflow.op.extract.split.constants import PARAGRAPH_SPLITTER
+
+
+class TestExtractImageFlow(unittest.TestCase):
+    @patch("uniflow.flow.extract.extract_image_flow.ExtractImageOp")
+    @patch("uniflow.flow.extract.extract_image_flow.ProcessImageOp")
+    @patch("uniflow.flow.extract.extract_image_flow.SplitterOpsFactory")
+    @patch("uniflow.flow.extract.extract_image_flow.CvModel")
+    def setUp(
+        self,
+        mock_cv_model,
+        mock_splitter_ops_factory,
+        mock_process_image_op,
+        mock_extract_image_op,
+    ):
+        self.mock_extract_image_op = mock_extract_image_op
+        self.mock_process_image_op = mock_process_image_op
+        self.mock_splitter_ops_factory = mock_splitter_ops_factory
+        self.mock_cv_model = mock_cv_model
+        self.mock_cv_model.return_value = MagicMock()
+        self.model_config = {"model_config": "model_config"}
+        self.extract_image_flow = ExtractImageFlow(model_config=self.model_config)
+
+    def test_init(self):
+        self.mock_extract_image_op.assert_called_once_with(
+            name="extract_image_op", model=self.mock_cv_model.return_value
+        )
+        self.mock_cv_model.assert_called_once_with(model_config=self.model_config)
+        self.mock_process_image_op.assert_called_once_with(name="process_image_op")
+        self.mock_splitter_ops_factory.get.assert_called_once_with(PARAGRAPH_SPLITTER)
+
+    def test_run(self):
+        # arrange
+        nodes = [
+            Node(name="node1", value_dict={"filename": "filepath"}),
+            Node(name="node2", value_dict={"filename": "filepath"}),
+        ]
+
+        self.mock_extract_image_op.return_value.return_value = MagicMock()
+        self.mock_process_image_op.return_value.return_value = MagicMock()
+        self.mock_splitter_ops_factory.get.return_value.return_value = MagicMock()
+        result = self.extract_image_flow.run(nodes)
+
+        self.mock_extract_image_op.return_value.assert_called_once_with(nodes)
+        self.mock_process_image_op.return_value.assert_called_once_with(
+            self.mock_extract_image_op.return_value.return_value
+        )
+        self.mock_splitter_ops_factory.get.return_value.assert_called_once_with(
+            self.mock_process_image_op.return_value.return_value
+        )
+
+        self.assertEqual(
+            result, self.mock_splitter_ops_factory.get.return_value.return_value
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/flow/extract/test_extract_ipynb_flow.py

@@ -0,0 +1,39 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from uniflow.flow.extract.extract_ipynb_flow import ExtractIpynbFlow
+from uniflow.node import Node
+
+
+class TestExtractIpynbFlow(unittest.TestCase):
+    @patch("uniflow.flow.extract.extract_ipynb_flow.ExtractIpynbOp")
+    @patch("uniflow.flow.extract.extract_ipynb_flow.ProcessIpynbOp")
+    def setUp(self, mock_process_ipynb_op, mock_extract_ipynb_op):
+        self.mock_process_ipynb_op = mock_process_ipynb_op
+        self.mock_extract_ipynb_op = mock_extract_ipynb_op
+        self.extract_ipynb_flow = ExtractIpynbFlow()
+
+    def test_init(self):
+        self.mock_extract_ipynb_op.assert_called_once_with(name="extract_ipynb_op")
+        self.mock_process_ipynb_op.assert_called_once_with(name="process_ipynb_op")
+
+    def test_run(self):
+        # arrange
+        nodes = [
+            Node(name="node1", value_dict={"filename": "filepath"}),
+            Node(name="node2", value_dict={"filename": "filepath"}),
+        ]
+
+        self.mock_process_ipynb_op.return_value.return_value = MagicMock()
+        self.mock_extract_ipynb_op.return_value.return_value = MagicMock()
+        result = self.extract_ipynb_flow.run(nodes)
+
+        self.mock_extract_ipynb_op.return_value.assert_called_once_with(nodes)
+        self.mock_process_ipynb_op.return_value.assert_called_once_with(
+            self.mock_extract_ipynb_op.return_value.return_value
+        )
+        self.assertEqual(result, self.mock_process_ipynb_op.return_value.return_value)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/flow/extract/test_extract_md_flow.py

@@ -0,0 +1,44 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from uniflow.flow.extract.extract_md_flow import ExtractMarkdownFlow
+from uniflow.node import Node
+from uniflow.op.extract.split.constants import MARKDOWN_HEADER_SPLITTER
+
+
+class TestExtractMarkdownFlow(unittest.TestCase):
+    @patch("uniflow.flow.extract.extract_md_flow.ExtractTxtOp")
+    @patch("uniflow.flow.extract.extract_md_flow.SplitterOpsFactory")
+    def setUp(self, mock_splitter_ops_factory, mock_extract_md_op):
+        self.mock_splitter_ops_factory = mock_splitter_ops_factory
+        self.mock_extract_md_op = mock_extract_md_op
+        self.extract_md_flow = ExtractMarkdownFlow()
+
+    def test_init(self):
+        self.mock_extract_md_op.assert_called_once_with(name="extract_md_op")
+        self.mock_splitter_ops_factory.get.assert_called_once_with(
+            MARKDOWN_HEADER_SPLITTER
+        )
+
+    def test_run(self):
+        # arrange
+        nodes = [
+            Node(name="node1", value_dict={"filename": "filepath"}),
+            Node(name="node2", value_dict={"filename": "filepath"}),
+        ]
+
+        self.mock_splitter_ops_factory.get.return_value.return_value = MagicMock()
+        self.mock_extract_md_op.return_value.return_value = MagicMock()
+        result = self.extract_md_flow.run(nodes)
+
+        self.mock_extract_md_op.return_value.assert_called_once_with(nodes)
+        self.mock_splitter_ops_factory.get.return_value.assert_called_once_with(
+            self.mock_extract_md_op.return_value.return_value
+        )
+        self.assertEqual(
+            result, self.mock_splitter_ops_factory.get.return_value.return_value
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/flow/extract/test_extract_pdf_flow.py

@@ -0,0 +1,63 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from uniflow.flow.extract.extract_pdf_flow import ExtractPDFFlow
+from uniflow.node import Node
+from uniflow.op.extract.split.constants import PARAGRAPH_SPLITTER
+
+
+class TestExtractPDFFlow(unittest.TestCase):
+    @patch("uniflow.flow.extract.extract_pdf_flow.ExtractPDFOp")
+    @patch("uniflow.flow.extract.extract_pdf_flow.ProcessPDFOp")
+    @patch("uniflow.flow.extract.extract_pdf_flow.SplitterOpsFactory")
+    @patch("uniflow.flow.extract.extract_pdf_flow.CvModel")
+    def setUp(
+        self,
+        mock_cv_model,
+        mock_splitter_ops_factory,
+        mock_process_pdf_op,
+        mock_extract_pdf_op,
+    ):
+        self.mock_extract_pdf_op = mock_extract_pdf_op
+        self.mock_process_pdf_op = mock_process_pdf_op
+        self.mock_splitter_ops_factory = mock_splitter_ops_factory
+        self.mock_cv_model = mock_cv_model
+        self.mock_cv_model.return_value = MagicMock()
+        self.model_config = {"model_config": "model_config"}
+        self.extract_pdf_flow = ExtractPDFFlow(model_config=self.model_config)
+
+    def test_init(self):
+        self.mock_extract_pdf_op.assert_called_once_with(
+            name="extract_pdf_op", model=self.mock_cv_model.return_value
+        )
+        self.mock_cv_model.assert_called_once_with(model_config=self.model_config)
+        self.mock_process_pdf_op.assert_called_once_with(name="process_pdf_op")
+        self.mock_splitter_ops_factory.get.assert_called_once_with(PARAGRAPH_SPLITTER)
+
+    def test_run(self):
+        # arrange
+        nodes = [
+            Node(name="node1", value_dict={"filename": "filepath"}),
+            Node(name="node2", value_dict={"filename": "filepath"}),
+        ]
+
+        self.mock_extract_pdf_op.return_value.return_value = MagicMock()
+        self.mock_process_pdf_op.return_value.return_value = MagicMock()
+        self.mock_splitter_ops_factory.get.return_value.return_value = MagicMock()
+        result = self.extract_pdf_flow.run(nodes)
+
+        self.mock_extract_pdf_op.return_value.assert_called_once_with(nodes)
+        self.mock_process_pdf_op.return_value.assert_called_once_with(
+            self.mock_extract_pdf_op.return_value.return_value
+        )
+        self.mock_splitter_ops_factory.get.return_value.assert_called_once_with(
+            self.mock_process_pdf_op.return_value.return_value
+        )
+
+        self.assertEqual(
+            result, self.mock_splitter_ops_factory.get.return_value.return_value
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/flow/extract/test_extract_txt_flow.py

@@ -1,89 +1,42 @@
 import unittest
-from unittest.mock import call, patch
+from unittest.mock import MagicMock, patch
 
 from uniflow.flow.extract.extract_txt_flow import ExtractTxtFlow
 from uniflow.node import Node
-from uniflow.op.extract.load.txt_op import ExtractTxtOp
 from uniflow.op.extract.split.constants import PARAGRAPH_SPLITTER
-from uniflow.op.extract.split.splitter_factory import SplitterOpsFactory
 
 
 class TestExtractTxtFlow(unittest.TestCase):
-
-    def setUp(self):
+    @patch("uniflow.flow.extract.extract_txt_flow.ExtractTxtOp")
+    @patch("uniflow.flow.extract.extract_txt_flow.SplitterOpsFactory")
+    def setUp(self, mock_splitter_ops_factory, mock_extract_txt_op):
+        self.mock_splitter_ops_factory = mock_splitter_ops_factory
+        self.mock_extract_txt_op = mock_extract_txt_op
         self.extract_txt_flow = ExtractTxtFlow()
-        self.extract_txt_flow._extract_txt_op = ExtractTxtOp(name="extract_txt_op")
-        self.extract_txt_flow._split_op = SplitterOpsFactory.get(PARAGRAPH_SPLITTER)
-
-    @patch(
-        "uniflow.op.extract.load.txt_op.read_file",
-        return_value="mocked\n\nfile\n\ncontent",
-    )
-    def test_call_with_empty_node(self, mock_read_file):
-        # arrange
-        nodes = []
 
-        # act
-        output_nodes = self.extract_txt_flow.run(nodes)
+    def test_init(self):
+        self.mock_extract_txt_op.assert_called_once_with(name="extract_txt_op")
+        self.mock_splitter_ops_factory.get.assert_called_once_with(PARAGRAPH_SPLITTER)
 
-        # assert
-        mock_read_file.assert_not_called()
-        self.assertEqual(len(output_nodes), 0)
-
-    @patch(
-        "uniflow.op.extract.load.txt_op.read_file",
-        return_value="mocked\n\nfile\n\ncontent",
-    )
-    def test_call_with_node_without_filename(self, mock_read_file):
+    def test_run(self):
         # arrange
-        node = Node(name="node1", value_dict={})
-
-        # act
-        with self.assertRaises(KeyError):
-            self.extract_txt_flow.run([node])
-
-        # assert
-        mock_read_file.assert_not_called()
-
-    @patch(
-        "uniflow.op.extract.load.txt_op.read_file",
-        return_value="mocked\n\nfile\n\ncontent",
-    )
-    def test_call_with_node(self, mock_read_file):
-        # arrange
-        node = Node(name="node1", value_dict={"filename": "mocked_file_path"})
-
-        # act
-        output_nodes = self.extract_txt_flow.run([node])
-
-        # assert
-        mock_read_file.assert_called_once_with("mocked_file_path")
-        self.assertEqual(len(output_nodes), 1)
+        nodes = [
+            Node(name="node1", value_dict={"filename": "filepath"}),
+            Node(name="node2", value_dict={"filename": "filepath"}),
+        ]
+
+        self.mock_splitter_ops_factory.get.return_value.return_value = MagicMock()
+        self.mock_extract_txt_op.return_value.return_value = MagicMock()
+        result = self.extract_txt_flow.run(nodes)
+
+        self.mock_extract_txt_op.return_value.assert_called_once_with(nodes)
+        self.mock_splitter_ops_factory.get.return_value.assert_called_once_with(
+            self.mock_extract_txt_op.return_value.return_value
+        )
         self.assertEqual(
-            output_nodes[0].value_dict["text"], ["mocked", "file", "content"]
+            result, self.mock_splitter_ops_factory.get.return_value.return_value
         )
 
-    @patch(
-        "uniflow.op.extract.load.txt_op.read_file",
-        return_value="mocked\n\nfile\n\ncontent",
-    )
-    def test_call_with_multiple_nodes(self, mock_read_file):
-        # arrange
-        node1 = Node(name="node1", value_dict={"filename": "mocked_file_path1"})
-        node2 = Node(name="node2", value_dict={"filename": "mocked_file_path2"})
-        nodes = [node1, node2]
 
-        # act
-        output_nodes = self.extract_txt_flow.run(nodes)
-
-        # assert
-        mock_read_file.assert_has_calls(
-            [call("mocked_file_path1"), call("mocked_file_path2")], any_order=True
-        )
-        self.assertEqual(len(output_nodes), 2)
-        self.assertEqual(
-            output_nodes[0].value_dict["text"], ["mocked", "file", "content"]
-        )
-        self.assertEqual(
-            output_nodes[1].value_dict["text"], ["mocked", "file", "content"]
-        )
+if __name__ == "__main__":
+    unittest.main()