merge feat/plugins

This commit is contained in:
Joel
2024-11-01 12:37:42 +08:00
653 changed files with 46039 additions and 3369 deletions

View File

@ -63,17 +63,24 @@ def test_run_invalid_variable_type(document_extractor_node, mock_graph_runtime_s
@pytest.mark.parametrize(
("mime_type", "file_content", "expected_text", "transfer_method"),
("mime_type", "file_content", "expected_text", "transfer_method", "extension"),
[
("text/plain", b"Hello, world!", ["Hello, world!"], FileTransferMethod.LOCAL_FILE),
("application/pdf", b"%PDF-1.5\n%Test PDF content", ["Mocked PDF content"], FileTransferMethod.LOCAL_FILE),
("text/plain", b"Hello, world!", ["Hello, world!"], FileTransferMethod.LOCAL_FILE, ".txt"),
(
"application/pdf",
b"%PDF-1.5\n%Test PDF content",
["Mocked PDF content"],
FileTransferMethod.LOCAL_FILE,
".pdf",
),
(
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
b"PK\x03\x04",
["Mocked DOCX content"],
FileTransferMethod.LOCAL_FILE,
FileTransferMethod.REMOTE_URL,
"",
),
("text/plain", b"Remote content", ["Remote content"], FileTransferMethod.REMOTE_URL),
("text/plain", b"Remote content", ["Remote content"], FileTransferMethod.REMOTE_URL, None),
],
)
def test_run_extract_text(
@ -83,6 +90,7 @@ def test_run_extract_text(
file_content,
expected_text,
transfer_method,
extension,
monkeypatch,
):
document_extractor_node.graph_runtime_state = mock_graph_runtime_state
@ -92,6 +100,7 @@ def test_run_extract_text(
mock_file.transfer_method = transfer_method
mock_file.related_id = "test_file_id" if transfer_method == FileTransferMethod.LOCAL_FILE else None
mock_file.remote_url = "https://example.com/file.txt" if transfer_method == FileTransferMethod.REMOTE_URL else None
mock_file.extension = extension
mock_array_file_segment = Mock(spec=ArrayFileSegment)
mock_array_file_segment.value = [mock_file]
@ -116,7 +125,7 @@ def test_run_extract_text(
result = document_extractor_node._run()
assert isinstance(result, NodeRunResult)
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED, result.error
assert result.outputs is not None
assert result.outputs["text"] == expected_text

View File

@ -192,7 +192,7 @@ def test_http_request_node_form_with_file(monkeypatch):
def attr_checker(*args, **kwargs):
assert kwargs["data"] == {"name": "test"}
assert kwargs["files"] == {"file": b"test"}
assert kwargs["files"] == {"file": (None, b"test", "application/octet-stream")}
return httpx.Response(200, content=b"")
monkeypatch.setattr(

View File

@ -55,6 +55,7 @@ def test_execute_if_else_result_true():
pool.add(["start", "less_than"], 21)
pool.add(["start", "greater_than_or_equal"], 22)
pool.add(["start", "less_than_or_equal"], 21)
pool.add(["start", "null"], None)
pool.add(["start", "not_null"], "1212")
node = IfElseNode(

View File

@ -33,8 +33,8 @@ def test_get_file_attribute(pool, file):
assert result.value == file.filename
# Test getting a non-existent attribute
with pytest.raises(ValueError):
pool.get(("node_1", "file_var", "non_existent_attr"))
result = pool.get(("node_1", "file_var", "non_existent_attr"))
assert result is None
def test_use_long_selector(pool):

View File

@ -0,0 +1,100 @@
import os
import posixpath
from unittest.mock import MagicMock
import pytest
from _pytest.monkeypatch import MonkeyPatch
from oss2 import Bucket
from oss2.models import GetObjectResult, PutObjectResult
from tests.unit_tests.oss.__mock.base import (
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
get_example_folder,
)
class MockResponse:
def __init__(self, status, headers, request_id):
self.status = status
self.headers = headers
self.request_id = request_id
class MockAliyunOssClass:
def __init__(
self,
auth,
endpoint,
bucket_name,
is_cname=False,
session=None,
connect_timeout=None,
app_name="",
enable_crc=True,
proxies=None,
region=None,
cloudbox_id=None,
is_path_style=False,
is_verify_object_strict=True,
):
self.bucket_name = get_example_bucket()
self.key = posixpath.join(get_example_folder(), get_example_filename())
self.content = get_example_data()
self.filepath = get_example_filepath()
self.resp = MockResponse(
200,
{
"etag": "ee8de918d05640145b18f70f4c3aa602",
"x-oss-version-id": "CAEQNhiBgMDJgZCA0BYiIDc4MGZjZGI2OTBjOTRmNTE5NmU5NmFhZjhjYmY0****",
},
"request_id",
)
def put_object(self, key, data, headers=None, progress_callback=None):
assert key == self.key
assert data == self.content
return PutObjectResult(self.resp)
def get_object(self, key, byte_range=None, headers=None, progress_callback=None, process=None, params=None):
assert key == self.key
get_object_output = MagicMock(GetObjectResult)
get_object_output.read.return_value = self.content
return get_object_output
def get_object_to_file(
self, key, filename, byte_range=None, headers=None, progress_callback=None, process=None, params=None
):
assert key == self.key
assert filename == self.filepath
def object_exists(self, key, headers=None):
assert key == self.key
return True
def delete_object(self, key, params=None, headers=None):
assert key == self.key
self.resp.headers["x-oss-delete-marker"] = True
return self.resp
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_aliyun_oss_mock(monkeypatch: MonkeyPatch):
if MOCK:
monkeypatch.setattr(Bucket, "__init__", MockAliyunOssClass.__init__)
monkeypatch.setattr(Bucket, "put_object", MockAliyunOssClass.put_object)
monkeypatch.setattr(Bucket, "get_object", MockAliyunOssClass.get_object)
monkeypatch.setattr(Bucket, "get_object_to_file", MockAliyunOssClass.get_object_to_file)
monkeypatch.setattr(Bucket, "object_exists", MockAliyunOssClass.object_exists)
monkeypatch.setattr(Bucket, "delete_object", MockAliyunOssClass.delete_object)
yield
if MOCK:
monkeypatch.undo()

View File

@ -0,0 +1,58 @@
from collections.abc import Generator
import pytest
from extensions.storage.base_storage import BaseStorage
def get_example_folder() -> str:
return "/dify"
def get_example_bucket() -> str:
return "dify"
def get_example_filename() -> str:
return "test.txt"
def get_example_data() -> bytes:
return b"test"
def get_example_filepath() -> str:
return "/test"
class BaseStorageTest:
@pytest.fixture(autouse=True)
def setup_method(self):
"""Should be implemented in child classes to setup specific storage."""
self.storage = BaseStorage()
def test_save(self):
"""Test saving data."""
self.storage.save(get_example_filename(), get_example_data())
def test_load_once(self):
"""Test loading data once."""
assert self.storage.load_once(get_example_filename()) == get_example_data()
def test_load_stream(self):
"""Test loading data as a stream."""
generator = self.storage.load_stream(get_example_filename())
assert isinstance(generator, Generator)
assert next(generator) == get_example_data()
def test_download(self):
"""Test downloading data."""
self.storage.download(get_example_filename(), get_example_filepath())
def test_exists(self):
"""Test checking if a file exists."""
assert self.storage.exists(get_example_filename())
def test_delete(self):
"""Test deleting a file."""
self.storage.delete(get_example_filename())

View File

@ -0,0 +1,57 @@
import os
import shutil
from pathlib import Path
from unittest.mock import MagicMock, mock_open, patch
import pytest
from _pytest.monkeypatch import MonkeyPatch
from tests.unit_tests.oss.__mock.base import (
get_example_data,
get_example_filename,
get_example_filepath,
get_example_folder,
)
class MockLocalFSClass:
def write_bytes(self, data):
assert data == get_example_data()
def read_bytes(self):
return get_example_data()
@staticmethod
def copyfile(src, dst):
assert src == os.path.join(get_example_folder(), get_example_filename())
assert dst == get_example_filepath()
@staticmethod
def exists(path):
assert path == os.path.join(get_example_folder(), get_example_filename())
return True
@staticmethod
def remove(path):
assert path == os.path.join(get_example_folder(), get_example_filename())
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_local_fs_mock(monkeypatch: MonkeyPatch):
if MOCK:
monkeypatch.setattr(Path, "write_bytes", MockLocalFSClass.write_bytes)
monkeypatch.setattr(Path, "read_bytes", MockLocalFSClass.read_bytes)
monkeypatch.setattr(shutil, "copyfile", MockLocalFSClass.copyfile)
monkeypatch.setattr(os.path, "exists", MockLocalFSClass.exists)
monkeypatch.setattr(os, "remove", MockLocalFSClass.remove)
os.makedirs = MagicMock()
with patch("builtins.open", mock_open(read_data=get_example_data())):
yield
if MOCK:
monkeypatch.undo()

View File

@ -0,0 +1,81 @@
import os
from unittest.mock import MagicMock
import pytest
from _pytest.monkeypatch import MonkeyPatch
from qcloud_cos import CosS3Client
from qcloud_cos.streambody import StreamBody
from tests.unit_tests.oss.__mock.base import (
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
)
class MockTencentCosClass:
def __init__(self, conf, retry=1, session=None):
self.bucket_name = get_example_bucket()
self.key = get_example_filename()
self.content = get_example_data()
self.filepath = get_example_filepath()
self.resp = {
"ETag": "ee8de918d05640145b18f70f4c3aa602",
"Server": "tencent-cos",
"x-cos-hash-crc64ecma": 16749565679157681890,
"x-cos-request-id": "NWU5MDNkYzlfNjRiODJhMDlfMzFmYzhfMTFm****",
}
def put_object(self, Bucket, Body, Key, EnableMD5=False, **kwargs): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
assert Body == self.content
return self.resp
def get_object(self, Bucket, Key, KeySimplifyCheck=True, **kwargs): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
mock_stream_body = MagicMock(StreamBody)
mock_raw_stream = MagicMock()
mock_stream_body.get_raw_stream.return_value = mock_raw_stream
mock_raw_stream.read.return_value = self.content
mock_stream_body.get_stream_to_file = MagicMock()
def chunk_generator(chunk_size=2):
for i in range(0, len(self.content), chunk_size):
yield self.content[i : i + chunk_size]
mock_stream_body.get_stream.return_value = chunk_generator(chunk_size=4096)
return {"Body": mock_stream_body}
def object_exists(self, Bucket, Key): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
return True
def delete_object(self, Bucket, Key, **kwargs): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
self.resp.update({"x-cos-delete-marker": True})
return self.resp
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_tencent_cos_mock(monkeypatch: MonkeyPatch):
if MOCK:
monkeypatch.setattr(CosS3Client, "__init__", MockTencentCosClass.__init__)
monkeypatch.setattr(CosS3Client, "put_object", MockTencentCosClass.put_object)
monkeypatch.setattr(CosS3Client, "get_object", MockTencentCosClass.get_object)
monkeypatch.setattr(CosS3Client, "object_exists", MockTencentCosClass.object_exists)
monkeypatch.setattr(CosS3Client, "delete_object", MockTencentCosClass.delete_object)
yield
if MOCK:
monkeypatch.undo()

View File

@ -1,5 +1,4 @@
import os
from typing import Union
from unittest.mock import MagicMock
import pytest
@ -7,28 +6,19 @@ from _pytest.monkeypatch import MonkeyPatch
from tos import TosClientV2
from tos.clientv2 import DeleteObjectOutput, GetObjectOutput, HeadObjectOutput, PutObjectOutput
from tests.unit_tests.oss.__mock.base import (
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
)
class AttrDict(dict):
def __getattr__(self, item):
return self.get(item)
def get_example_bucket() -> str:
return "dify"
def get_example_filename() -> str:
return "test.txt"
def get_example_data() -> bytes:
return b"test"
def get_example_filepath() -> str:
return "/test"
class MockVolcengineTosClass:
def __init__(self, ak="", sk="", endpoint="", region=""):
self.bucket_name = get_example_bucket()

View File

@ -0,0 +1,22 @@
from unittest.mock import MagicMock, patch
import pytest
from oss2 import Auth
from extensions.storage.aliyun_oss_storage import AliyunOssStorage
from tests.unit_tests.oss.__mock.aliyun_oss import setup_aliyun_oss_mock
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_bucket,
get_example_folder,
)
class TestAliyunOss(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_aliyun_oss_mock):
"""Executed before each test method."""
with patch.object(Auth, "__init__", return_value=None):
self.storage = AliyunOssStorage()
self.storage.bucket_name = get_example_bucket()
self.storage.folder = get_example_folder()

View File

@ -0,0 +1,18 @@
from collections.abc import Generator
import pytest
from extensions.storage.local_fs_storage import LocalFsStorage
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_folder,
)
from tests.unit_tests.oss.__mock.local import setup_local_fs_mock
class TestLocalFS(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_local_fs_mock):
"""Executed before each test method."""
self.storage = LocalFsStorage()
self.storage.folder = get_example_folder()

View File

@ -0,0 +1,20 @@
from unittest.mock import patch
import pytest
from qcloud_cos import CosConfig
from extensions.storage.tencent_cos_storage import TencentCosStorage
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_bucket,
)
from tests.unit_tests.oss.__mock.tencent_cos import setup_tencent_cos_mock
class TestTencentCos(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_tencent_cos_mock):
"""Executed before each test method."""
with patch.object(CosConfig, "__init__", return_value=None):
self.storage = TencentCosStorage()
self.storage.bucket_name = get_example_bucket()

View File

@ -1,30 +1,18 @@
from collections.abc import Generator
from flask import Flask
import pytest
from tos import TosClientV2
from tos.clientv2 import GetObjectOutput, HeadObjectOutput, PutObjectOutput
from extensions.storage.volcengine_tos_storage import VolcengineTosStorage
from tests.unit_tests.oss.__mock.volcengine_tos import (
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
setup_volcengine_tos_mock,
)
from tests.unit_tests.oss.__mock.volcengine_tos import setup_volcengine_tos_mock
class VolcengineTosTest:
_instance = None
def __new__(cls):
if cls._instance == None:
cls._instance = object.__new__(cls)
return cls._instance
else:
return cls._instance
def __init__(self):
class TestVolcengineTos(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_volcengine_tos_mock):
"""Executed before each test method."""
self.storage = VolcengineTosStorage()
self.storage.bucket_name = get_example_bucket()
self.storage.client = TosClientV2(
@ -33,35 +21,3 @@ class VolcengineTosTest:
endpoint="https://xxx.volces.com",
region="cn-beijing",
)
def test_save(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
volc_tos.storage.save(get_example_filename(), get_example_data())
def test_load_once(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
assert volc_tos.storage.load_once(get_example_filename()) == get_example_data()
def test_load_stream(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
generator = volc_tos.storage.load_stream(get_example_filename())
assert isinstance(generator, Generator)
assert next(generator) == get_example_data()
def test_download(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
volc_tos.storage.download(get_example_filename(), get_example_filepath())
def test_exists(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
assert volc_tos.storage.exists(get_example_filename())
def test_delete(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
volc_tos.storage.delete(get_example_filename())