diff --git a/api/tests/unit_tests/oss/__mock/tencent_cos.py b/api/tests/unit_tests/oss/__mock/tencent_cos.py new file mode 100644 index 0000000000..5189b68e87 --- /dev/null +++ b/api/tests/unit_tests/oss/__mock/tencent_cos.py @@ -0,0 +1,81 @@ +import os +from unittest.mock import MagicMock + +import pytest +from _pytest.monkeypatch import MonkeyPatch +from qcloud_cos import CosS3Client +from qcloud_cos.streambody import StreamBody + +from tests.unit_tests.oss.__mock.base import ( + get_example_bucket, + get_example_data, + get_example_filename, + get_example_filepath, +) + + +class MockTencentCosClass: + def __init__(self, conf, retry=1, session=None): + self.bucket_name = get_example_bucket() + self.key = get_example_filename() + self.content = get_example_data() + self.filepath = get_example_filepath() + self.resp = { + "ETag": "ee8de918d05640145b18f70f4c3aa602", + "Server": "tencent-cos", + "x-cos-hash-crc64ecma": 16749565679157681890, + "x-cos-request-id": "NWU5MDNkYzlfNjRiODJhMDlfMzFmYzhfMTFm****", + } + + def put_object(self, Bucket, Body, Key, EnableMD5=False, **kwargs): # noqa: N803 + assert Bucket == self.bucket_name + assert Key == self.key + assert Body == self.content + return self.resp + + def get_object(self, Bucket, Key, KeySimplifyCheck=True, **kwargs): # noqa: N803 + assert Bucket == self.bucket_name + assert Key == self.key + + mock_stream_body = MagicMock(StreamBody) + mock_raw_stream = MagicMock() + mock_stream_body.get_raw_stream.return_value = mock_raw_stream + mock_raw_stream.read.return_value = self.content + + mock_stream_body.get_stream_to_file = MagicMock() + + def chunk_generator(chunk_size=2): + for i in range(0, len(self.content), chunk_size): + yield self.content[i : i + chunk_size] + + mock_stream_body.get_stream.return_value = chunk_generator(chunk_size=4096) + return {"Body": mock_stream_body} + + def object_exists(self, Bucket, Key): # noqa: N803 + assert Bucket == self.bucket_name + assert Key == self.key + return True + + def delete_object(self, Bucket, Key, **kwargs): # noqa: N803 + assert Bucket == self.bucket_name + assert Key == self.key + self.resp.update({"x-cos-delete-marker": True}) + return self.resp + + +MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true" + + +@pytest.fixture +def setup_tencent_cos_mock(monkeypatch: MonkeyPatch): + if MOCK: + monkeypatch.setattr(CosS3Client, "__init__", MockTencentCosClass.__init__) + monkeypatch.setattr(CosS3Client, "put_object", MockTencentCosClass.put_object) + monkeypatch.setattr(CosS3Client, "get_object", MockTencentCosClass.get_object) + monkeypatch.setattr(CosS3Client, "object_exists", MockTencentCosClass.object_exists) + monkeypatch.setattr(CosS3Client, "delete_object", MockTencentCosClass.delete_object) + + yield + + if MOCK: + monkeypatch.undo() diff --git a/api/tests/unit_tests/oss/tencent_cos/__init__.py b/api/tests/unit_tests/oss/tencent_cos/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py b/api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py new file mode 100644 index 0000000000..303f0493bd --- /dev/null +++ b/api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py @@ -0,0 +1,20 @@ +from unittest.mock import patch + +import pytest +from qcloud_cos import CosConfig + +from extensions.storage.tencent_cos_storage import TencentCosStorage +from tests.unit_tests.oss.__mock.base import ( + BaseStorageTest, + get_example_bucket, +) +from tests.unit_tests.oss.__mock.tencent_cos import setup_tencent_cos_mock + + +class TestTencentCos(BaseStorageTest): + @pytest.fixture(autouse=True) + def setup_method(self, setup_tencent_cos_mock): + """Executed before each test method.""" + with patch.object(CosConfig, "__init__", return_value=None): + self.storage = TencentCosStorage() + self.storage.bucket_name = get_example_bucket() 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 index 9f8aa158a9..5afbc9e8b4 100644 --- a/api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py +++ b/api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py @@ -1,5 +1,3 @@ -from collections.abc import Generator - import pytest from tos import TosClientV2