| @@ -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()) | |||
| @@ -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() | |||
| @@ -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() | |||
| @@ -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() | |||
| @@ -1,30 +1,20 @@ | |||
| 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 +23,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()) | |||