From b2aa385942a74f7cba5cfa2e32e9027ca5860bd5 Mon Sep 17 00:00:00 2001 From: ice yao Date: Fri, 18 Oct 2024 20:08:41 +0800 Subject: [PATCH] feat: Add volcengine tos storage test (#9495) --- api/tests/unit_tests/oss/__mock/__init__.py | 0 .../unit_tests/oss/__mock/volcengine_tos.py | 100 ++++++++++++++++++ .../unit_tests/oss/volcengine_tos/__init__.py | 0 .../oss/volcengine_tos/test_volcengine_tos.py | 67 ++++++++++++ 4 files changed, 167 insertions(+) create mode 100644 api/tests/unit_tests/oss/__mock/__init__.py create mode 100644 api/tests/unit_tests/oss/__mock/volcengine_tos.py create mode 100644 api/tests/unit_tests/oss/volcengine_tos/__init__.py create mode 100644 api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py diff --git a/api/tests/unit_tests/oss/__mock/__init__.py b/api/tests/unit_tests/oss/__mock/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/api/tests/unit_tests/oss/__mock/volcengine_tos.py b/api/tests/unit_tests/oss/__mock/volcengine_tos.py new file mode 100644 index 0000000000..241764c521 --- /dev/null +++ b/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() diff --git a/api/tests/unit_tests/oss/volcengine_tos/__init__.py b/api/tests/unit_tests/oss/volcengine_tos/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py b/api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py new file mode 100644 index 0000000000..3f334a3764 --- /dev/null +++ b/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())