test: use real object storage SDK responses (#43231)

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

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."""

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

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

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."""