|
| 1 | +"""Workflow Web sources persist images before splitting and bind them on write.""" |
| 2 | + |
| 3 | +import os |
| 4 | +from unittest import skipUnless |
| 5 | +from unittest.mock import MagicMock, patch |
| 6 | +from uuid import uuid4 |
| 7 | + |
| 8 | +import boto3 |
| 9 | +from botocore.config import Config |
| 10 | +from django.test import SimpleTestCase, TestCase |
| 11 | + |
| 12 | +from application.workflow.common import WorkflowType |
| 13 | +from application.workflow.nodes.data_source_web_node.data_source_web_node import DataSourceWebNode |
| 14 | +from application.workflow.nodes.document_split_node.document_split_node import DocumentSplitNode |
| 15 | +from application.workflow.nodes.knowledge_write_node.knowledge_write_node import KnowledgeWriteNode |
| 16 | +from common.utils.fork import ChildLink, Fork |
| 17 | +from knowledge.models import File, FileSourceType, Knowledge, KnowledgeType, Paragraph, ParagraphAsset |
| 18 | + |
| 19 | + |
| 20 | +NODE_MODULE = "application.workflow.nodes.data_source_web_node.data_source_web_node" |
| 21 | + |
| 22 | + |
| 23 | +def source_node(parameters): |
| 24 | + node = DataSourceWebNode.__new__(DataSourceWebNode) |
| 25 | + node.get_workflow_parameters = MagicMock(return_value=parameters) |
| 26 | + node._check_cancelled = MagicMock() |
| 27 | + node.get_node_id = MagicMock(return_value="web") |
| 28 | + node.write_context = MagicMock() |
| 29 | + return node |
| 30 | + |
| 31 | + |
| 32 | +class WorkflowWebImageSourceTests(SimpleTestCase): |
| 33 | + def test_full_import_and_single_document_sync_localize_images_before_output(self): |
| 34 | + for single in (False, True): |
| 35 | + with self.subTest(single=single): |
| 36 | + parameters = { |
| 37 | + "knowledge_id": "knowledge-id", |
| 38 | + "data_source": {"source_url": "https://example.com/page.html"}, |
| 39 | + "workflow_source": {"source_scope": "web-scope"}, |
| 40 | + } |
| 41 | + if single: |
| 42 | + parameters["workflow_sync_document"] = {"name": "existing title"} |
| 43 | + node = source_node(parameters) |
| 44 | + remote = "" |
| 45 | + local = "" |
| 46 | + response = Fork.Response.success(remote, []) |
| 47 | + |
| 48 | + def crawl(_depth, _visited, handler): |
| 49 | + handler(ChildLink(parameters["data_source"]["source_url"], None), response) |
| 50 | + |
| 51 | + with ( |
| 52 | + patch(f"{NODE_MODULE}.Fork") as fetch, |
| 53 | + patch(f"{NODE_MODULE}.ForkManage") as crawler, |
| 54 | + patch(f"{NODE_MODULE}.internalize_web_images", return_value=local) as internalize, |
| 55 | + ): |
| 56 | + fetch.return_value.fork.return_value = response |
| 57 | + crawler.return_value.fork.side_effect = crawl |
| 58 | + node.execute() |
| 59 | + |
| 60 | + internalize.assert_called_once_with(remote, "knowledge-id") |
| 61 | + document = node.write_context.call_args_list[0].args[1][0] |
| 62 | + self.assertEqual(document["content"], local) |
| 63 | + self.assertEqual(document["meta"]["source_type"], "web") |
| 64 | + self.assertEqual(document["meta"]["source_scope"], "web-scope") |
| 65 | + self.assertEqual(document["meta"]["source_url"], parameters["data_source"]["source_url"]) |
| 66 | + if single: |
| 67 | + self.assertEqual(document["name"], "existing title") |
| 68 | + |
| 69 | + def test_failed_page_does_not_download_images_or_publish_a_document(self): |
| 70 | + documents = [] |
| 71 | + handler = source_node({})._get_collect_handler(documents, knowledge_id="knowledge-id") |
| 72 | + with patch(f"{NODE_MODULE}.internalize_web_images") as internalize: |
| 73 | + with self.assertRaisesRegex(ValueError, "unavailable"): |
| 74 | + handler(ChildLink("https://example.com/page.html", None), Fork.Response.error("unavailable")) |
| 75 | + internalize.assert_not_called() |
| 76 | + self.assertEqual(documents, []) |
| 77 | + |
| 78 | + |
| 79 | +@skipUnless(os.environ.get("WORKFLOW_WEB_IMAGE_TEST_URL"), "Start the local workflow Web Postman simulation") |
| 80 | +class WorkflowWebImageHTTPTests(TestCase): |
| 81 | + def setUp(self): |
| 82 | + self.base_url = os.environ["WORKFLOW_WEB_IMAGE_TEST_URL"].rstrip("/") |
| 83 | + self.knowledge = Knowledge.objects.create(name="Workflow Web images", desc="", type=KnowledgeType.WORKFLOW) |
| 84 | + self.parameters = { |
| 85 | + "knowledge_id": str(self.knowledge.id), |
| 86 | + "knowledge_action_id": str(uuid4()), |
| 87 | + "workflow_source": {"source_scope": "web-scope", "source_type": "web"}, |
| 88 | + "data_source": {"source_url": f"{self.base_url}/site/page.html"}, |
| 89 | + } |
| 90 | + client = boto3.client( |
| 91 | + "s3", |
| 92 | + endpoint_url=f"{self.base_url}/storage", |
| 93 | + aws_access_key_id="test-key", |
| 94 | + aws_secret_access_key="test-secret", |
| 95 | + region_name="us-east-1", |
| 96 | + config=Config(signature_version="s3v4", s3={"addressing_style": "path"}), |
| 97 | + ) |
| 98 | + self.enterContext(patch("knowledge.models.knowledge.get_s3_client", return_value=client)) |
| 99 | + self.enterContext(patch("knowledge.models.knowledge.get_bucket", return_value="bucket")) |
| 100 | + self.enterContext(patch("knowledge.models.knowledge.is_seaweedfs_enabled", return_value=True)) |
| 101 | + |
| 102 | + def import_page(self): |
| 103 | + web = source_node(self.parameters) |
| 104 | + web.execute() |
| 105 | + downloaded = web.write_context.call_args_list[0].args[1] |
| 106 | + splitter = DocumentSplitNode.__new__(DocumentSplitNode) |
| 107 | + splitter.get_parameters = MagicMock(return_value={"document_list": ["web", "document_list"]}) |
| 108 | + splitter.get_workflow_parameters = MagicMock(return_value=self.parameters) |
| 109 | + splitter.get_workflow_type = MagicMock(return_value=WorkflowType.KNOWLEDGE) |
| 110 | + splitter.get_reference_content = MagicMock(return_value=downloaded) |
| 111 | + splitter.write_context = MagicMock() |
| 112 | + splitter.execute() |
| 113 | + documents = next( |
| 114 | + call.args[1] for call in splitter.write_context.call_args_list if call.args[0] == "paragraph_list" |
| 115 | + ) |
| 116 | + writer = KnowledgeWriteNode.__new__(KnowledgeWriteNode) |
| 117 | + writer.get_workflow_parameters = MagicMock(return_value=self.parameters) |
| 118 | + return writer.save(documents, None)[0][0] |
| 119 | + |
| 120 | + def test_crawl_download_split_and_write_share_one_local_image(self): |
| 121 | + document = self.import_page() |
| 122 | + file = File.objects.get() |
| 123 | + self.assertEqual(file.source_type, FileSourceType.PARAGRAPH) |
| 124 | + self.assertTrue(Paragraph.objects.filter(id=file.source_id, document=document).exists()) |
| 125 | + assets = list(ParagraphAsset.objects.filter(document=document)) |
| 126 | + self.assertEqual(len(assets), 2) |
| 127 | + self.assertEqual({asset.file_id for asset in assets}, {file.id}) |
| 128 | + self.assertTrue(file.get_bytes().startswith(b"\x89PNG\r\n\x1a\n")) |
| 129 | + paragraphs = list(Paragraph.objects.filter(document=document)) |
| 130 | + self.assertFalse(any(f"{self.base_url}/site/chart.png" in p.content for p in paragraphs)) |
| 131 | + self.assertTrue(any(f"/oss/file/{file.id}" in p.content for p in paragraphs)) |
| 132 | + self.assertEqual(document.meta["source_url"], self.parameters["data_source"]["source_url"]) |
| 133 | + second = self.import_page() |
| 134 | + self.assertEqual(File.objects.count(), 1) |
| 135 | + self.assertEqual( |
| 136 | + set(ParagraphAsset.objects.filter(document=second).values_list("file_id", flat=True)), {file.id} |
| 137 | + ) |
| 138 | + |
| 139 | + def test_missing_image_preserves_remote_reference_and_imports_text(self): |
| 140 | + self.parameters["data_source"]["source_url"] = f"{self.base_url}/site/missing.html" |
| 141 | + self.parameters["workflow_sync_document"] = {"name": "retained title"} |
| 142 | + document = self.import_page() |
| 143 | + self.assertEqual(document.name, "retained title") |
| 144 | + self.assertFalse(File.objects.exists()) |
| 145 | + self.assertFalse(ParagraphAsset.objects.exists()) |
| 146 | + self.assertTrue(Paragraph.objects.filter(document=document, content__contains="/site/missing.png").exists()) |
0 commit comments