瀏覽代碼

chore: add local storage test (#9827)

ice yao 6 月之前
父節點
當前提交
fc2297a2ca

+ 58 - 0
api/tests/unit_tests/oss/__mock/base.py

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

+ 57 - 0
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()

+ 7 - 17
api/tests/unit_tests/oss/__mock/volcengine_tos.py

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

+ 18 - 0
api/tests/unit_tests/oss/local/test_local_fs.py

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

+ 8 - 50
api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py

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