feat: Add volcengine tos storage test (#9495)

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

@@ -0,0 +1,100 @@
+import os
+from typing import Union
+from unittest.mock import MagicMock
+
+import pytest
+from _pytest.monkeypatch import MonkeyPatch
+from tos import TosClientV2
+from tos.clientv2 import DeleteObjectOutput, GetObjectOutput, HeadObjectOutput, PutObjectOutput
+
+
+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()
+        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()

api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py

@@ -0,0 +1,67 @@
+from collections.abc import Generator
+
+from flask import Flask
+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 (
+    get_example_bucket,
+    get_example_data,
+    get_example_filename,
+    get_example_filepath,
+    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):
+        self.storage = VolcengineTosStorage(app=Flask(__name__))
+        self.storage.bucket_name = get_example_bucket()
+        self.storage.client = TosClientV2(
+            ak="dify",
+            sk="dify",
+            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())