This commit is contained in:
Joel
2024-10-28 16:33:02 +08:00
parent 966d42a4db
commit 0cec6195a3
9 changed files with 42 additions and 84 deletions

View File

@ -65,8 +65,7 @@ def test_run_invalid_variable_type(document_extractor_node, mock_graph_runtime_s
@pytest.mark.parametrize(
("mime_type", "file_content", "expected_text", "transfer_method", "extension"),
[
("text/plain", b"Hello, world!",
["Hello, world!"], FileTransferMethod.LOCAL_FILE, ".txt"),
("text/plain", b"Hello, world!", ["Hello, world!"], FileTransferMethod.LOCAL_FILE, ".txt"),
(
"application/pdf",
b"%PDF-1.5\n%Test PDF content",
@ -81,8 +80,7 @@ def test_run_invalid_variable_type(document_extractor_node, mock_graph_runtime_s
FileTransferMethod.REMOTE_URL,
"",
),
("text/plain", b"Remote content",
["Remote content"], FileTransferMethod.REMOTE_URL, None),
("text/plain", b"Remote content", ["Remote content"], FileTransferMethod.REMOTE_URL, None),
],
)
def test_run_extract_text(
@ -119,12 +117,10 @@ def test_run_extract_text(
if mime_type == "application/pdf":
mock_pdf_extract = Mock(return_value=expected_text[0])
monkeypatch.setattr(
"core.workflow.nodes.document_extractor.node._extract_text_from_pdf", mock_pdf_extract)
monkeypatch.setattr("core.workflow.nodes.document_extractor.node._extract_text_from_pdf", mock_pdf_extract)
elif mime_type.startswith("application/vnd.openxmlformats"):
mock_docx_extract = Mock(return_value=expected_text[0])
monkeypatch.setattr(
"core.workflow.nodes.document_extractor.node._extract_text_from_doc", mock_docx_extract)
monkeypatch.setattr("core.workflow.nodes.document_extractor.node._extract_text_from_doc", mock_docx_extract)
result = document_extractor_node._run()
@ -134,8 +130,7 @@ def test_run_extract_text(
assert result.outputs["text"] == expected_text
if transfer_method == FileTransferMethod.REMOTE_URL:
mock_ssrf_proxy_get.assert_called_once_with(
"https://example.com/file.txt")
mock_ssrf_proxy_get.assert_called_once_with("https://example.com/file.txt")
elif transfer_method == FileTransferMethod.LOCAL_FILE:
mock_download.assert_called_once_with(mock_file)

View File

@ -77,18 +77,12 @@ MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_volcengine_tos_mock(monkeypatch: MonkeyPatch):
if MOCK:
monkeypatch.setattr(TosClientV2, "__init__",
MockVolcengineTosClass.__init__)
monkeypatch.setattr(TosClientV2, "put_object",
MockVolcengineTosClass.put_object)
monkeypatch.setattr(TosClientV2, "get_object",
MockVolcengineTosClass.get_object)
monkeypatch.setattr(TosClientV2, "get_object_to_file",
MockVolcengineTosClass.get_object_to_file)
monkeypatch.setattr(TosClientV2, "head_object",
MockVolcengineTosClass.head_object)
monkeypatch.setattr(TosClientV2, "delete_object",
MockVolcengineTosClass.delete_object)
monkeypatch.setattr(TosClientV2, "__init__", MockVolcengineTosClass.__init__)
monkeypatch.setattr(TosClientV2, "put_object", MockVolcengineTosClass.put_object)
monkeypatch.setattr(TosClientV2, "get_object", MockVolcengineTosClass.get_object)
monkeypatch.setattr(TosClientV2, "get_object_to_file", MockVolcengineTosClass.get_object_to_file)
monkeypatch.setattr(TosClientV2, "head_object", MockVolcengineTosClass.head_object)
monkeypatch.setattr(TosClientV2, "delete_object", MockVolcengineTosClass.delete_object)
yield