Browse Source

chore: add local storage test (#9827)

tags/0.10.2
ice yao 1 year ago
parent
commit
fc2297a2ca
No account linked to committer's email address

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

@@ -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 View File

@@ -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 View File

@@ -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 View File

@@ -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 View File

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

Loading…
Cancel
Save