Merge pull request #198 from jli943/main

add test_extract_flow

Author: SayaZhang

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

@@ -0,0 +1,42 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from uniflow.flow.extract.extract_txt_flow import ExtractTxtFlow
+from uniflow.node import Node
+from uniflow.op.extract.split.constants import PARAGRAPH_SPLITTER
+
+
+class TestExtractTxtFlow(unittest.TestCase):
+    @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()
+
+    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)
+
+    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_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(
+            result, self.mock_splitter_ops_factory.get.return_value.return_value
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()