diff --git a/api/tests/unit_tests/oss/__mock/aliyun_oss.py b/api/tests/unit_tests/oss/__mock/aliyun_oss.py index 59434e072fa..6138990c884 100644 --- a/api/tests/unit_tests/oss/__mock/aliyun_oss.py +++ b/api/tests/unit_tests/oss/__mock/aliyun_oss.py @@ -1,9 +1,9 @@ import os import posixpath -from unittest.mock import MagicMock import pytest from oss2 import Bucket +from oss2.http import Response as OssResponse from oss2.models import GetObjectResult, PutObjectResult from tests.unit_tests.oss.__mock.base import ( @@ -12,6 +12,7 @@ from tests.unit_tests.oss.__mock.base import ( get_example_filename, get_example_filepath, get_example_folder, + object_response, ) @@ -60,9 +61,9 @@ class MockAliyunOssClass: def get_object(self, key, byte_range=None, headers=None, progress_callback=None, process=None, params=None): assert key == self.key - get_object_output = MagicMock(GetObjectResult) - get_object_output.read.return_value = self.content - return get_object_output + response = object_response(self.content) + response.headers.update(self.resp.headers) + return GetObjectResult(OssResponse(response)) def get_object_to_file( self, key, filename, byte_range=None, headers=None, progress_callback=None, process=None, params=None diff --git a/api/tests/unit_tests/oss/__mock/base.py b/api/tests/unit_tests/oss/__mock/base.py index 5bde461d944..dfcdc70d9fd 100644 --- a/api/tests/unit_tests/oss/__mock/base.py +++ b/api/tests/unit_tests/oss/__mock/base.py @@ -1,10 +1,20 @@ from collections.abc import Generator +from io import BytesIO import pytest +from requests import Response from extensions.storage.base_storage import BaseStorage +def object_response(content: bytes) -> Response: + response = Response() + response.status_code = 200 + response.headers["Content-Length"] = str(len(content)) + response.raw = BytesIO(content) + return response + + def get_example_folder() -> str: return "~/dify" @@ -53,6 +63,7 @@ class BaseStorageTest: generator = self.storage.load_stream(get_example_filename()) assert isinstance(generator, Generator) assert next(generator) == get_example_data() + assert list(generator) == [] def test_download(self): """Test downloading data.""" diff --git a/api/tests/unit_tests/oss/__mock/tencent_cos.py b/api/tests/unit_tests/oss/__mock/tencent_cos.py index 65b5bf31329..40fe96a0d78 100644 --- a/api/tests/unit_tests/oss/__mock/tencent_cos.py +++ b/api/tests/unit_tests/oss/__mock/tencent_cos.py @@ -1,5 +1,4 @@ import os -from unittest.mock import MagicMock import pytest from qcloud_cos import CosS3Client @@ -10,6 +9,7 @@ from tests.unit_tests.oss.__mock.base import ( get_example_data, get_example_filename, get_example_filepath, + object_response, ) @@ -36,19 +36,7 @@ class MockTencentCosClass: 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} + return {"Body": StreamBody(object_response(self.content))} def object_exists(self, Bucket, Key): # noqa: N803 assert Bucket == self.bucket_name diff --git a/api/tests/unit_tests/oss/__mock/volcengine_tos.py b/api/tests/unit_tests/oss/__mock/volcengine_tos.py index af35fdeb2d0..6b9a9199422 100644 --- a/api/tests/unit_tests/oss/__mock/volcengine_tos.py +++ b/api/tests/unit_tests/oss/__mock/volcengine_tos.py @@ -1,16 +1,17 @@ import os from collections import UserDict -from unittest.mock import MagicMock import pytest from tos import TosClientV2 from tos.clientv2 import DeleteObjectOutput, GetObjectOutput, HeadObjectOutput, PutObjectOutput +from tos.http import Response as TosResponse from tests.unit_tests.oss.__mock.base import ( get_example_bucket, get_example_data, get_example_filename, get_example_filepath, + object_response, ) @@ -51,9 +52,7 @@ class MockVolcengineTosClass: 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 + return GetObjectOutput(TosResponse(object_response(self.content))) def get_object_to_file(self, bucket: str, key: str, file_path: str): assert bucket == self.bucket_name 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 index 2802a2f1e39..76a88e8a116 100644 --- a/api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py +++ b/api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py @@ -1,3 +1,4 @@ +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -7,6 +8,8 @@ from extensions.storage.tencent_cos_storage import TencentCosStorage from tests.unit_tests.oss.__mock.base import ( BaseStorageTest, get_example_bucket, + get_example_data, + get_example_filename, ) pytest_plugins = ("tests.unit_tests.oss.__mock.tencent_cos",) @@ -20,6 +23,11 @@ class TestTencentCos(BaseStorageTest): self.storage = TencentCosStorage() self.storage.bucket_name = get_example_bucket() + def test_download(self, tmp_path: Path): + destination = tmp_path / "download.txt" + self.storage.download(get_example_filename(), str(destination)) + assert destination.read_bytes() == get_example_data() + class TestTencentCosConfiguration: """Tests for TencentCosStorage initialization with different configurations."""