File size: 2,809 Bytes
a8b3f00 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
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()
|