dify

dify/api/tests/unit_tests/oss/__mock/aliyun_oss.py

@@ -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()

dify/api/tests/unit_tests/oss/__mock/base.py

@@ -0,0 +1,65 @@
+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_opendal_bucket() -> str:
+    return "./dify"
+
+
+def get_example_filename() -> str:
+    return "test.txt"
+
+
+def get_example_data(length: int = 4) -> bytes:
+    chars = "test"
+    result = "".join(chars[i % len(chars)] for i in range(length)).encode()
+    assert len(result) == length
+    return result
+
+
+def get_example_filepath() -> str:
+    return "~/test"
+
+
+class BaseStorageTest:
+    @pytest.fixture(autouse=True)
+    def setup_method(self, *args, **kwargs):
+        """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())

dify/api/tests/unit_tests/oss/__mock/local.py

@@ -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()

dify/api/tests/unit_tests/oss/__mock/tencent_cos.py

@@ -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()

dify/api/tests/unit_tests/oss/__mock/volcengine_tos.py

@@ -0,0 +1,91 @@
+import os
+from collections import UserDict
+from unittest.mock import MagicMock
+
+import pytest
+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(UserDict):
+    def __getattr__(self, item):
+        return self.get(item)
+
+
+class MockVolcengineTosClass:
+    def __init__(self, ak="", sk="", endpoint="", region=""):
+        self.bucket_name = get_example_bucket()
+        self.key = get_example_filename()
+        self.content = get_example_data()
+        self.filepath = get_example_filepath()
+        self.resp = AttrDict(
+            {
+                "x-tos-server-side-encryption": "kms",
+                "x-tos-server-side-encryption-kms-key-id": "trn:kms:cn-beijing:****:keyrings/ring-test/keys/key-test",
+                "x-tos-server-side-encryption-customer-algorithm": "AES256",
+                "x-tos-version-id": "test",
+                "x-tos-hash-crc64ecma": 123456,
+                "request_id": "test",
+                "headers": {
+                    "x-tos-id-2": "test",
+                    "ETag": "123456",
+                },
+                "status": 200,
+            }
+        )
+
+    def put_object(self, bucket: str, key: str, content=None) -> PutObjectOutput:
+        assert bucket == self.bucket_name
+        assert key == self.key
+        assert content == self.content
+        return PutObjectOutput(self.resp)
+
+    def get_object(self, bucket: str, key: str) -> GetObjectOutput:
+        assert bucket == self.bucket_name
+        assert key == self.key
+
+        get_object_output = MagicMock(GetObjectOutput)
+        get_object_output.read.return_value = self.content
+        return get_object_output
+
+    def get_object_to_file(self, bucket: str, key: str, file_path: str):
+        assert bucket == self.bucket_name
+        assert key == self.key
+        assert file_path == self.filepath
+
+    def head_object(self, bucket: str, key: str) -> HeadObjectOutput:
+        assert bucket == self.bucket_name
+        assert key == self.key
+        return HeadObjectOutput(self.resp)
+
+    def delete_object(self, bucket: str, key: str):
+        assert bucket == self.bucket_name
+        assert key == self.key
+        return DeleteObjectOutput(self.resp)
+
+
+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)
+
+    yield
+
+    if MOCK:
+        monkeypatch.undo()