test(storage): P3-1第十二波 SharedStorageService单测 55个 #716

Merged
xiaoxia merged 4 commits from feat/p3-1-storage-tests into develop 2026-07-22 16:26:56 +08:00
3 changed files with 1429 additions and 252 deletions
+337
View File
@@ -0,0 +1,337 @@
"""
pagination 通用分页器单元测试
覆盖:
- PaginationParams: 默认值/边界/校验/offset/limit
- PaginationMeta: from_params 各种边界场景
- PaginatedResponse: create 工厂方法
- paginate: 内存分页函数
"""
import pytest
from pydantic import ValidationError
from packages.application.common.pagination import (
PaginatedResponse,
PaginationMeta,
PaginationParams,
paginate,
)
# ============================================================
# PaginationParams
# ============================================================
class TestPaginationParamsDefaults:
"""默认值测试"""
def test_default_page_is_1(self):
params = PaginationParams()
assert params.page == 1
def test_default_page_size_is_20(self):
params = PaginationParams()
assert params.page_size == 20
def test_default_offset_is_0(self):
params = PaginationParams()
assert params.offset == 0
def test_default_limit_is_20(self):
params = PaginationParams()
assert params.limit == 20
class TestPaginationParamsValidation:
"""参数校验"""
@pytest.mark.parametrize("page", [1, 2, 100, 9999])
def test_valid_page_values(self, page):
params = PaginationParams(page=page)
assert params.page == page
def test_page_zero_raises(self):
with pytest.raises(ValidationError):
PaginationParams(page=0)
def test_page_negative_raises(self):
with pytest.raises(ValidationError):
PaginationParams(page=-1)
@pytest.mark.parametrize("page_size", [1, 20, 50, 100])
def test_valid_page_size_values(self, page_size):
params = PaginationParams(page_size=page_size)
assert params.page_size == page_size
def test_page_size_zero_raises(self):
with pytest.raises(ValidationError):
PaginationParams(page_size=0)
def test_page_size_negative_raises(self):
with pytest.raises(ValidationError):
PaginationParams(page_size=-5)
def test_page_size_over_100_raises(self):
with pytest.raises(ValidationError):
PaginationParams(page_size=101)
def test_invalid_page_type_raises(self):
with pytest.raises(ValidationError):
PaginationParams(page="abc")
def test_invalid_page_size_type_raises(self):
with pytest.raises(ValidationError):
PaginationParams(page_size="abc")
class TestPaginationParamsOffset:
"""offset 属性计算"""
def test_page_1_offset_0(self):
params = PaginationParams(page=1, page_size=20)
assert params.offset == 0
def test_page_2_offset_page_size(self):
params = PaginationParams(page=2, page_size=20)
assert params.offset == 20
def test_page_3_offset_2x_page_size(self):
params = PaginationParams(page=3, page_size=20)
assert params.offset == 40
def test_page_5_page_size_10_offset_40(self):
params = PaginationParams(page=5, page_size=10)
assert params.offset == 40
def test_page_1_page_size_100_offset_0(self):
params = PaginationParams(page=1, page_size=100)
assert params.offset == 0
class TestPaginationParamsLimit:
"""limit 属性"""
def test_limit_equals_page_size(self):
params = PaginationParams(page_size=20)
assert params.limit == 20
def test_limit_1(self):
params = PaginationParams(page_size=1)
assert params.limit == 1
def test_limit_100(self):
params = PaginationParams(page_size=100)
assert params.limit == 100
# ============================================================
# PaginationMeta.from_params
# ============================================================
class TestPaginationMetaFromParams:
"""from_params 工厂方法"""
def test_empty_total_zero(self):
params = PaginationParams(page=1, page_size=20)
meta = PaginationMeta.from_params(params, total=0)
assert meta.total == 0
assert meta.total_pages == 0
assert meta.has_next is False
assert meta.has_prev is False
def test_exactly_one_page(self):
params = PaginationParams(page=1, page_size=20)
meta = PaginationMeta.from_params(params, total=20)
assert meta.total_pages == 1
assert meta.has_next is False
assert meta.has_prev is False
def test_less_than_one_page(self):
params = PaginationParams(page=1, page_size=20)
meta = PaginationMeta.from_params(params, total=15)
assert meta.total_pages == 1
assert meta.has_next is False
assert meta.has_prev is False
def test_multiple_pages_first_page(self):
params = PaginationParams(page=1, page_size=20)
meta = PaginationMeta.from_params(params, total=50)
assert meta.total_pages == 3
assert meta.has_next is True
assert meta.has_prev is False
def test_multiple_pages_middle_page(self):
params = PaginationParams(page=2, page_size=20)
meta = PaginationMeta.from_params(params, total=50)
assert meta.total_pages == 3
assert meta.has_next is True
assert meta.has_prev is True
def test_multiple_pages_last_page(self):
params = PaginationParams(page=3, page_size=20)
meta = PaginationMeta.from_params(params, total=50)
assert meta.total_pages == 3
assert meta.has_next is False
assert meta.has_prev is True
def test_exact_division(self):
params = PaginationParams(page=2, page_size=20)
meta = PaginationMeta.from_params(params, total=40)
assert meta.total_pages == 2
assert meta.has_next is False
assert meta.has_prev is True
def test_non_exact_division_ceil(self):
params = PaginationParams(page=1, page_size=20)
meta = PaginationMeta.from_params(params, total=41)
assert meta.total_pages == 3
def test_total_1_page_size_20(self):
params = PaginationParams(page=1, page_size=20)
meta = PaginationMeta.from_params(params, total=1)
assert meta.total_pages == 1
assert meta.has_next is False
assert meta.has_prev is False
def test_page_beyond_total_pages(self):
params = PaginationParams(page=10, page_size=20)
meta = PaginationMeta.from_params(params, total=50)
assert meta.total_pages == 3
assert meta.has_next is False
assert meta.has_prev is True
def test_preserves_params_values(self):
params = PaginationParams(page=3, page_size=15)
meta = PaginationMeta.from_params(params, total=100)
assert meta.page == 3
assert meta.page_size == 15
assert meta.total == 100
# ============================================================
# PaginatedResponse.create
# ============================================================
class TestPaginatedResponseCreate:
"""create 工厂方法"""
def test_create_with_data(self):
params = PaginationParams(page=1, page_size=20)
data = [1, 2, 3]
response = PaginatedResponse.create(data, params, total=100)
assert response.data == data
assert response.pagination.total == 100
assert response.pagination.page == 1
assert response.pagination.page_size == 20
def test_create_with_empty_data(self):
params = PaginationParams(page=1, page_size=20)
response = PaginatedResponse.create([], params, total=0)
assert response.data == []
assert response.pagination.total == 0
assert response.pagination.total_pages == 0
def test_create_preserves_list_type(self):
params = PaginationParams(page=1, page_size=20)
data = ["a", "b", "c"]
response = PaginatedResponse.create(data, params, total=10)
assert response.data == ["a", "b", "c"]
assert len(response.data) == 3
# ============================================================
# paginate 函数
# ============================================================
class TestPaginateFunction:
"""内存分页函数"""
def test_empty_list(self):
params = PaginationParams(page=1, page_size=20)
result = paginate([], params)
assert result.data == []
assert result.pagination.total == 0
assert result.pagination.total_pages == 0
def test_first_page(self):
items = list(range(50))
params = PaginationParams(page=1, page_size=20)
result = paginate(items, params)
assert result.data == list(range(20))
assert result.pagination.total == 50
assert result.pagination.total_pages == 3
assert result.pagination.has_next is True
assert result.pagination.has_prev is False
def test_middle_page(self):
items = list(range(50))
params = PaginationParams(page=2, page_size=20)
result = paginate(items, params)
assert result.data == list(range(20, 40))
assert result.pagination.has_next is True
assert result.pagination.has_prev is True
def test_last_page(self):
items = list(range(50))
params = PaginationParams(page=3, page_size=20)
result = paginate(items, params)
assert result.data == list(range(40, 50))
assert len(result.data) == 10
assert result.pagination.has_next is False
assert result.pagination.has_prev is True
def test_page_beyond_total(self):
items = list(range(25))
params = PaginationParams(page=10, page_size=20)
result = paginate(items, params)
assert result.data == []
assert result.pagination.total == 25
assert result.pagination.total_pages == 2
def test_page_size_larger_than_total(self):
items = list(range(5))
params = PaginationParams(page=1, page_size=20)
result = paginate(items, params)
assert result.data == items
assert result.pagination.total_pages == 1
assert result.pagination.has_next is False
def test_single_item(self):
items = [42]
params = PaginationParams(page=1, page_size=20)
result = paginate(items, params)
assert result.data == [42]
assert result.pagination.total == 1
def test_page_size_1(self):
items = list(range(5))
params = PaginationParams(page=3, page_size=1)
result = paginate(items, params)
assert result.data == [2]
assert result.pagination.total_pages == 5
def test_exact_page_size(self):
items = list(range(40))
params = PaginationParams(page=2, page_size=20)
result = paginate(items, params)
assert result.data == list(range(20, 40))
assert result.pagination.total_pages == 2
assert result.pagination.has_next is False
def test_string_items(self):
items = ["a", "b", "c", "d", "e"]
params = PaginationParams(page=2, page_size=2)
result = paginate(items, params)
assert result.data == ["c", "d"]
assert result.pagination.total == 5
def test_does_not_mutate_original_list(self):
items = list(range(10))
original = items.copy()
params = PaginationParams(page=1, page_size=3)
paginate(items, params)
assert items == original
+539
View File
@@ -0,0 +1,539 @@
"""
SharedStorageService 单元测试
重点覆盖纯逻辑部分:
- _normalize_storage_key: URL提取 + URL解码
- _is_local_generated_url: 本地生成URL判断
- get_url: 公共URL拼接
- create_direct_upload_post: policy + HMAC签名
- get_download_url: bucket=None时的fallback
- 未配置OSS时的错误处理
- 单例模式
"""
import base64
import hashlib
import hmac
import json
import os
from unittest.mock import MagicMock, patch
import pytest
from packages.shared.storage import (
SharedStorageService,
get_shared_storage_service,
get_storage_service,
)
# ============================================================
# Fixtures
# ============================================================
def _make_service(
bucket_name="test-bucket",
endpoint="oss-cn-hangzhou.aliyuncs.com",
access_key_id="test-key-id",
access_key_secret="test-key-secret",
local_url_prefix="/generated-files",
with_bucket=True,
):
"""创建一个 SharedStorageService 实例,mock 掉 oss2 和 settings。"""
mock_settings = MagicMock()
mock_settings.oss_bucket_name = bucket_name
mock_settings.oss_endpoint = endpoint
mock_settings.oss_access_key_id = access_key_id
mock_settings.oss_access_key_secret = access_key_secret
mock_bucket = MagicMock() if with_bucket else None
with (
patch("packages.shared.storage.get_shared_settings", return_value=mock_settings),
patch.dict(os.environ, {"GENERATED_FILES_URL_PREFIX": local_url_prefix}, clear=False),
):
if with_bucket:
with patch("packages.shared.storage.oss2") as mock_oss2:
mock_oss2.Auth.return_value = MagicMock()
mock_oss2.Bucket.return_value = mock_bucket
service = SharedStorageService()
service.bucket = mock_bucket
return service, mock_bucket, mock_settings
else:
service = SharedStorageService()
service.bucket = None
return service, None, mock_settings
# ============================================================
# _normalize_storage_key
# ============================================================
class TestNormalizeStorageKey:
"""_normalize_storage_key URL 提取与解码"""
def test_plain_key_returns_as_is(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("videos/clip.mp4")
assert result == "videos/clip.mp4"
def test_key_with_leading_slash_stripped(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("/videos/clip.mp4")
assert result == "videos/clip.mp4"
def test_https_url_extracts_path(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://test-bucket.oss-cn-hangzhou.aliyuncs.com/videos/clip.mp4")
assert result == "videos/clip.mp4"
def test_http_url_extracts_path(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("http://test-bucket.oss-cn-hangzhou.aliyuncs.com/audio/voice.mp3")
assert result == "audio/voice.mp3"
def test_url_with_query_strips_query(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss-cn.com/file.mp4?signature=abc&expires=123")
assert result == "file.mp4"
def test_url_with_leading_slash_in_path(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss.com//double/slash.jpg")
assert result == "double/slash.jpg"
def test_url_decodes_percent_encoded_spaces(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss.com/my%20video.mp4")
assert result == "my video.mp4"
def test_url_decodes_percent_encoded_chinese(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss.com/%E4%B8%AD%E6%96%87.mp4")
assert result == "中文.mp4"
def test_url_with_special_chars_decoded(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss.com/file%281%29.jpg")
assert result == "file(1).jpg"
def test_plain_key_with_percent_not_decoded(self):
"""原始 key 不以 http 开头,不做 URL 解码,直接 lstrip('/')"""
service, _, _ = _make_service()
result = service._normalize_storage_key("file%20name.mp4")
# 不是 URL,直接返回(去掉前导/)
assert result == "file%20name.mp4"
def test_empty_string(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("")
assert result == ""
def test_root_slash_url(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss.com/")
assert result == ""
def test_nested_path_url(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss.com/a/b/c/d/file.txt")
assert result == "a/b/c/d/file.txt"
def test_url_with_port(self):
service, _, _ = _make_service()
result = service._normalize_storage_key("https://bucket.oss.com:443/file.txt")
assert result == "file.txt"
# ============================================================
# _is_local_generated_url
# ============================================================
class TestIsLocalGeneratedUrl:
"""_is_local_generated_url 本地URL判断"""
def test_local_prefix_returns_true(self):
service, _, _ = _make_service()
assert service._is_local_generated_url("/generated-files/abc.mp4") is True
def test_relative_local_returns_true(self):
service, _, _ = _make_service()
# 没有 scheme,直接用原字符串匹配
assert service._is_local_generated_url("/generated-files/out.mp4") is True
def test_full_url_with_local_path_returns_true(self):
service, _, _ = _make_service()
assert service._is_local_generated_url("https://example.com/generated-files/abc.mp4") is True
def test_other_path_returns_false(self):
service, _, _ = _make_service()
assert service._is_local_generated_url("/videos/abc.mp4") is False
def test_empty_string_returns_false(self):
service, _, _ = _make_service()
assert service._is_local_generated_url("") is False
def test_custom_prefix(self):
service, _, _ = _make_service(local_url_prefix="/custom-prefix")
assert service._is_local_generated_url("/custom-prefix/file.mp4") is True
assert service._is_local_generated_url("/generated-files/file.mp4") is False
# ============================================================
# get_url
# ============================================================
class TestGetUrl:
"""get_url 公共URL拼接"""
def test_returns_public_url_plus_key(self):
service, _, _ = _make_service()
result = service.get_url("videos/test.mp4")
assert result == "https://test-bucket.oss-cn-hangzhou.aliyuncs.com/videos/test.mp4"
def test_empty_key(self):
service, _, _ = _make_service()
result = service.get_url("")
assert result == "https://test-bucket.oss-cn-hangzhou.aliyuncs.com/"
def test_custom_bucket_and_endpoint(self):
service, _, _ = _make_service(
bucket_name="my-bucket",
endpoint="oss-us-east-1.aliyuncs.com",
)
result = service.get_url("file.txt")
assert result == "https://my-bucket.oss-us-east-1.aliyuncs.com/file.txt"
# ============================================================
# create_direct_upload_post
# ============================================================
class TestCreateDirectUploadPost:
"""create_direct_upload_post 直传表单生成"""
def test_returns_dict_with_expected_keys(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post(
storage_key="uploads/test.jpg",
content_type="image/jpeg",
max_size_bytes=10 * 1024 * 1024,
expires_seconds=3600,
)
assert "url" in result
assert "method" in result
assert "storage_key" in result
assert "expires_at" in result
assert "fields" in result
def test_method_is_post(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 1024, 3600)
assert result["method"] == "POST"
def test_url_is_public_url(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 1024, 3600)
assert result["url"] == "https://test-bucket.oss-cn-hangzhou.aliyuncs.com"
def test_storage_key_normalized(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post("/uploads/test.jpg", "image/jpeg", 1024, 3600)
assert result["storage_key"] == "uploads/test.jpg"
assert result["fields"]["key"] == "uploads/test.jpg"
def test_fields_contain_required_keys(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 1024, 3600)
fields = result["fields"]
assert fields["key"] == "uploads/a.jpg"
assert fields["OSSAccessKeyId"] == "test-key-id"
assert fields["success_action_status"] == "201"
assert fields["Content-Type"] == "image/jpeg"
assert "policy" in fields
assert "Signature" in fields
def test_policy_signature_is_valid_hmac_sha1(self):
"""验证 HMAC-SHA1 签名是否正确"""
secret = "my-secret-key-123"
service, _, _ = _make_service(access_key_secret=secret)
result = service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 1024, 3600)
policy = result["fields"]["policy"]
signature = result["fields"]["Signature"]
# 手动计算签名验证
expected = base64.b64encode(
hmac.new(secret.encode("utf-8"), policy.encode("utf-8"), hashlib.sha1).digest()
).decode("ascii")
assert signature == expected
def test_policy_contains_bucket_and_key(self):
service, _, _ = _make_service(bucket_name="my-bucket")
result = service.create_direct_upload_post("uploads/photo.png", "image/png", 2048, 1800)
policy = json.loads(base64.b64decode(result["fields"]["policy"]))
conditions = policy["conditions"]
assert {"bucket": "my-bucket"} in conditions
assert {"key": "uploads/photo.png"} in conditions
def test_policy_contains_content_length_range(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 5242880, 3600)
policy = json.loads(base64.b64decode(result["fields"]["policy"]))
conditions = policy["conditions"]
size_condition = [c for c in conditions if isinstance(c, list) and c[0] == "content-length-range"]
assert len(size_condition) == 1
assert size_condition[0][1] == 1
assert size_condition[0][2] == 5242880
def test_policy_content_type_starts_with(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 1024, 3600)
policy = json.loads(base64.b64decode(result["fields"]["policy"]))
conditions = policy["conditions"]
ct_condition = [c for c in conditions if isinstance(c, list) and c[0] == "starts-with"]
assert len(ct_condition) == 1
assert ct_condition[0][1] == "$Content-Type"
assert ct_condition[0][2] == "image/"
def test_policy_has_expiration(self):
service, _, _ = _make_service()
result = service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 1024, 3600)
policy = json.loads(base64.b64decode(result["fields"]["policy"]))
assert "expiration" in policy
# ISO 8601 格式
assert policy["expiration"].endswith("Z")
def test_non_uploads_key_raises_value_error(self):
service, _, _ = _make_service()
with pytest.raises(ValueError, match="uploads/"):
service.create_direct_upload_post("videos/a.mp4", "video/mp4", 1024, 3600)
def test_no_credentials_raises_runtime_error(self):
service, _, _ = _make_service(access_key_id="", access_key_secret="", with_bucket=False)
with pytest.raises(RuntimeError, match="not configured"):
service.create_direct_upload_post("uploads/a.jpg", "image/jpeg", 1024, 3600)
def test_url_normalized_key_in_uploads(self):
service, _, _ = _make_service()
# URL 形式的 key 被 normalize 后如果在 uploads/ 下应该可以
result = service.create_direct_upload_post(
"https://test-bucket.oss-cn-hangzhou.aliyuncs.com/uploads/from_url.jpg",
"image/jpeg",
1024,
3600,
)
assert result["storage_key"] == "uploads/from_url.jpg"
# ============================================================
# get_download_url (bucket=None 时的 fallback)
# ============================================================
class TestGetDownloadUrlFallback:
"""get_download_url 在 bucket 未配置时的 fallback 逻辑"""
def test_no_bucket_local_url_returns_as_is(self):
service, _, _ = _make_service(with_bucket=False)
result = service.get_download_url("/generated-files/test.mp4")
assert result == "/generated-files/test.mp4"
def test_no_bucket_regular_key_returns_public_url(self):
service, _, _ = _make_service(with_bucket=False)
result = service.get_download_url("videos/test.mp4")
assert result == "https://test-bucket.oss-cn-hangzhou.aliyuncs.com/videos/test.mp4"
def test_no_bucket_url_input_normalized(self):
service, _, _ = _make_service(with_bucket=False)
result = service.get_download_url("https://test-bucket.oss-cn-hangzhou.aliyuncs.com/videos/test.mp4")
assert result == "https://test-bucket.oss-cn-hangzhou.aliyuncs.com/videos/test.mp4"
def test_with_bucket_calls_sign_url(self):
service, mock_bucket, _ = _make_service(with_bucket=True)
mock_bucket.sign_url.return_value = "https://signed-url.com/file?sig=abc"
result = service.get_download_url("videos/test.mp4", expires_seconds=7200)
mock_bucket.sign_url.assert_called_once_with("GET", "videos/test.mp4", 7200)
assert result == "https://signed-url.com/file?sig=abc"
def test_sign_url_exception_falls_back_to_public_url(self):
service, mock_bucket, _ = _make_service(with_bucket=True)
mock_bucket.sign_url.side_effect = Exception("sign error")
result = service.get_download_url("videos/test.mp4")
assert result == "https://test-bucket.oss-cn-hangzhou.aliyuncs.com/videos/test.mp4"
# ============================================================
# 未配置 OSS 时的错误处理
# ============================================================
class TestNoBucketErrorHandling:
"""bucket=None 时的错误处理"""
def test_upload_file_raises(self):
service, _, _ = _make_service(with_bucket=False)
with pytest.raises(RuntimeError, match="not configured"):
service.upload_file("/tmp/test.txt", "uploads/test.txt")
def test_download_file_raises(self):
service, _, _ = _make_service(with_bucket=False)
with pytest.raises(RuntimeError, match="not configured"):
service.download_file("uploads/test.txt", "/tmp/test.txt")
def test_delete_file_silent_noop(self):
service, _, _ = _make_service(with_bucket=False)
# 不抛异常
result = service.delete_file("uploads/test.txt")
assert result is None
def test_file_exists_returns_false(self):
service, _, _ = _make_service(with_bucket=False)
assert service.file_exists("uploads/test.txt") is False
# ============================================================
# upload_file / delete_file / file_exists 正常路径
# ============================================================
class TestBucketOperations:
"""有 bucket 时的操作调用验证"""
def test_upload_file_with_path_string(self):
service, mock_bucket, _ = _make_service()
result = service.upload_file("/tmp/file.txt", "uploads/file.txt", "text/plain")
mock_bucket.put_object_from_file.assert_called_once()
args = mock_bucket.put_object_from_file.call_args
assert args[0][0] == "uploads/file.txt"
assert args[0][1] == "/tmp/file.txt"
assert result.startswith("https://test-bucket.")
def test_upload_file_with_file_object(self):
service, mock_bucket, _ = _make_service()
mock_file = MagicMock()
result = service.upload_file(mock_file, "uploads/file.bin", "application/octet-stream")
mock_file.seek.assert_called_once_with(0)
mock_bucket.put_object.assert_called_once()
assert result.startswith("https://test-bucket.")
def test_delete_file_calls_bucket(self):
service, mock_bucket, _ = _make_service()
service.delete_file("uploads/test.txt")
mock_bucket.delete_object.assert_called_once_with("uploads/test.txt")
def test_delete_file_exception_logged_not_raised(self):
service, mock_bucket, _ = _make_service()
mock_bucket.delete_object.side_effect = Exception("delete error")
# 不抛异常
service.delete_file("uploads/test.txt")
def test_file_exists_delegates_to_bucket(self):
service, mock_bucket, _ = _make_service()
mock_bucket.object_exists.return_value = True
assert service.file_exists("some/key") is True
mock_bucket.object_exists.assert_called_once_with("some/key")
def test_file_exists_false(self):
service, mock_bucket, _ = _make_service()
mock_bucket.object_exists.return_value = False
assert service.file_exists("some/key") is False
# ============================================================
# 单例 & 兼容别名
# ============================================================
class TestSingleton:
"""get_shared_storage_service 单例模式"""
def test_get_storage_service_is_alias(self):
# 两个函数返回同一个实例
with patch("packages.shared.storage._storage_service", None):
with patch("packages.shared.storage.SharedStorageService") as mock_cls:
mock_instance = MagicMock()
mock_cls.return_value = mock_instance
svc1 = get_shared_storage_service()
svc2 = get_storage_service()
assert svc1 is svc2
# 因为是同一个单例,类只实例化一次
assert mock_cls.call_count == 1
# ============================================================
# __init__ endpoint 处理
# ============================================================
class TestInitEndpointHandling:
"""初始化时 endpoint https 前缀处理"""
def test_endpoint_without_https_gets_prefix(self):
mock_settings = MagicMock()
mock_settings.oss_bucket_name = "test-bucket"
mock_settings.oss_endpoint = "oss-cn-hangzhou.aliyuncs.com"
mock_settings.oss_access_key_id = "key-id"
mock_settings.oss_access_key_secret = "key-secret"
with (
patch("packages.shared.storage.get_shared_settings", return_value=mock_settings),
patch("packages.shared.storage.oss2") as mock_oss2,
):
mock_oss2.Bucket.return_value = MagicMock()
service = SharedStorageService()
# 验证 Bucket 构造时 endpoint 带了 https://
call_args = mock_oss2.Bucket.call_args
assert call_args[0][1] == "https://oss-cn-hangzhou.aliyuncs.com"
def test_endpoint_with_https_keeps_as_is(self):
mock_settings = MagicMock()
mock_settings.oss_bucket_name = "test-bucket"
mock_settings.oss_endpoint = "https://oss-cn-hangzhou.aliyuncs.com"
mock_settings.oss_access_key_id = "key-id"
mock_settings.oss_access_key_secret = "key-secret"
with (
patch("packages.shared.storage.get_shared_settings", return_value=mock_settings),
patch("packages.shared.storage.oss2") as mock_oss2,
):
mock_oss2.Bucket.return_value = MagicMock()
service = SharedStorageService()
call_args = mock_oss2.Bucket.call_args
assert call_args[0][1] == "https://oss-cn-hangzhou.aliyuncs.com"
def test_endpoint_with_http_keeps_as_is(self):
mock_settings = MagicMock()
mock_settings.oss_bucket_name = "test-bucket"
mock_settings.oss_endpoint = "http://oss-cn-hangzhou.aliyuncs.com"
mock_settings.oss_access_key_id = "key-id"
mock_settings.oss_access_key_secret = "key-secret"
with (
patch("packages.shared.storage.get_shared_settings", return_value=mock_settings),
patch("packages.shared.storage.oss2") as mock_oss2,
):
mock_oss2.Bucket.return_value = MagicMock()
service = SharedStorageService()
call_args = mock_oss2.Bucket.call_args
assert call_args[0][1] == "http://oss-cn-hangzhou.aliyuncs.com"
+553 -252
View File
@@ -1,296 +1,597 @@
"""URL 安全校验工具单元测试 — SSRF 防护."""
"""
url_security URL安全校验单元测试
from __future__ import annotations
覆盖:
- validate_url_safety: scheme/主机/端口/SSRF/内网域名/白名单
- is_url_safe: 便捷函数
- UrlSecurityError / NoRedirectHandler
- _validate_magic_number: 文件魔数校验
- safe_download_file / safe_download_bytes: mock 网络测试
"""
import os
import sys
import unittest
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "apps", "worker"))
import shutil
import tempfile
from unittest.mock import MagicMock, patch
from video_processing.url_security import ( # noqa: E402
import pytest
from packages.shared.url_security import (
ALLOWED_AUDIO_MIME_TYPES,
ALLOWED_IMAGE_MIME_TYPES,
ALLOWED_PORTS,
ALLOWED_SCHEMES,
MAX_URL_LENGTH,
NoRedirectHandler,
UrlSecurityError,
_check_internal_hostnames,
_is_trusted_domain,
_validate_magic_number,
is_url_safe,
safe_download_bytes,
safe_download_file,
validate_url_safety,
)
class TestUrlSecurityValidation(unittest.TestCase):
"""URL 安全校验测试."""
# ── Scheme 白名单 ──────────────────────────────────────────────────────
def test_http_scheme_allowed(self):
"""HTTP scheme 应该被允许."""
result = validate_url_safety("http://example.com/test", purpose="test")
self.assertEqual(result, "http://example.com/test")
def test_https_scheme_allowed(self):
"""HTTPS scheme 应该被允许."""
result = validate_url_safety("https://example.com/test", purpose="test")
self.assertEqual(result, "https://example.com/test")
def test_file_scheme_rejected(self):
"""file:// scheme 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("file:///etc/passwd", purpose="test")
def test_ftp_scheme_rejected(self):
"""ftp:// scheme 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("ftp://example.com/test", purpose="test")
def test_empty_scheme_rejected(self):
"""空 scheme 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("example.com/test", purpose="test")
# ── 端口白名单 ────────────────────────────────────────────────────────
def test_port_80_allowed(self):
"""端口 80 应该被允许."""
# 80端口是默认HTTP端口,不显式指定也可以
result = validate_url_safety("http://example.com:80/test", purpose="test")
self.assertIn("example.com", result)
def test_port_443_allowed(self):
"""端口 443 应该被允许."""
result = validate_url_safety("https://example.com:443/test", purpose="test")
self.assertIn("example.com", result)
def test_port_8080_rejected(self):
"""非标准端口 8080 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://example.com:8080/test", purpose="test")
def test_port_22_rejected(self):
"""SSH 端口 22 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://example.com:22/test", purpose="test")
# ── SSRF: 直接 IP 访问 ───────────────────────────────────────────────
def test_loopback_ip_rejected(self):
"""回环地址 127.0.0.1 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://127.0.0.1/test", purpose="test")
def test_private_ip_192_rejected(self):
"""内网地址 192.168.x.x 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://192.168.1.1/test", purpose="test")
def test_private_ip_10_rejected(self):
"""内网地址 10.x.x.x 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://10.0.0.1/test", purpose="test")
def test_private_ip_172_rejected(self):
"""内网地址 172.16.x.x 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://172.16.0.1/test", purpose="test")
def test_unspecified_ip_rejected(self):
"""未指定地址 0.0.0.0 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://0.0.0.0/test", purpose="test")
def test_ipv6_loopback_rejected(self):
"""IPv6 回环 ::1 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://[::1]/test", purpose="test")
def test_ipv6_link_local_rejected(self):
"""IPv6 链路本地地址应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://[fe80::1]/test", purpose="test")
# ── SSRF: 内网主机名 ─────────────────────────────────────────────────
def test_localhost_rejected(self):
"""localhost 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://localhost/test", purpose="test")
def test_local_domain_rejected(self):
""".local 域名应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://printer.local/test", purpose="test")
def test_internal_domain_rejected(self):
""".internal 域名应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http://db.internal/test", purpose="test")
# ── URL 格式校验 ─────────────────────────────────────────────────────
def test_empty_url_rejected(self):
"""空 URL 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("", purpose="test")
def test_none_url_rejected(self):
"""None URL 应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety(None, purpose="test") # type: ignore
def test_url_too_long_rejected(self):
"""超长 URL 应该被拒绝."""
long_url = "https://example.com/" + "a" * 3000
with self.assertRaises(UrlSecurityError):
validate_url_safety(long_url, purpose="test")
def test_no_hostname_rejected(self):
"""缺少主机名应该被拒绝."""
with self.assertRaises(UrlSecurityError):
validate_url_safety("http:///test", purpose="test")
# ── is_url_safe 便捷函数 ─────────────────────────────────────────────
def test_is_url_safe_true(self):
"""安全 URL 应该返回 True."""
self.assertTrue(is_url_safe("https://example.com/test", purpose="test"))
def test_is_url_safe_false(self):
"""不安全 URL 应该返回 False."""
self.assertFalse(is_url_safe("http://127.0.0.1/test", purpose="test"))
def test_is_url_safe_empty(self):
"""空 URL 应该返回 False."""
self.assertFalse(is_url_safe("", purpose="test"))
# ── validate_url_safety 基础校验 ─────────────────────────────────────────────
if __name__ == "__main__":
unittest.main()
class TestValidateUrlSafetyBasics:
"""URL 安全校验基础测试"""
def test_valid_http_url(self):
url = "http://example.com/file.mp4"
result = validate_url_safety(url)
assert result == url
def test_valid_https_url(self):
url = "https://example.com/file.mp4"
result = validate_url_safety(url)
assert result == url
def test_empty_url_raises(self):
with pytest.raises(UrlSecurityError, match="为空"):
validate_url_safety("")
def test_none_url_raises(self):
with pytest.raises(UrlSecurityError):
validate_url_safety(None)
def test_url_too_long_raises(self):
long_url = "https://example.com/" + "a" * 2050
with pytest.raises(UrlSecurityError, match="过长"):
validate_url_safety(long_url)
def test_url_at_max_length_ok(self):
base = "https://example.com/"
pad = "a" * (MAX_URL_LENGTH - len(base))
url = base + pad
assert len(url) <= MAX_URL_LENGTH
result = validate_url_safety(url)
assert result == url
def test_invalid_scheme_ftp_raises(self):
with pytest.raises(UrlSecurityError, match="scheme"):
validate_url_safety("ftp://example.com/file")
def test_invalid_scheme_file_raises(self):
with pytest.raises(UrlSecurityError, match="scheme"):
validate_url_safety("file:///etc/passwd")
def test_invalid_scheme_data_raises(self):
with pytest.raises(UrlSecurityError, match="scheme"):
validate_url_safety("data:text/html,<script>")
def test_missing_scheme_raises(self):
with pytest.raises(UrlSecurityError, match="scheme"):
validate_url_safety("example.com/file")
def test_missing_hostname_raises(self):
with pytest.raises(UrlSecurityError, match="主机名"):
validate_url_safety("http:///path")
def test_uppercase_scheme_normalized(self):
"""HTTP/HTTPS 大写也能通过"""
url = "HTTPS://example.com/file"
# scheme 检查用 lower 比较
result = validate_url_safety(url)
assert result == url
def test_default_port_80_ok(self):
url = "http://example.com:80/file"
result = validate_url_safety(url)
assert result == url
def test_default_port_443_ok(self):
url = "https://example.com:443/file"
result = validate_url_safety(url)
assert result == url
def test_non_standard_port_raises(self):
with pytest.raises(UrlSecurityError, match="端口"):
validate_url_safety("http://example.com:8080/file")
def test_port_22_ssh_raises(self):
with pytest.raises(UrlSecurityError, match="端口"):
validate_url_safety("http://example.com:22/file")
def test_port_3306_mysql_raises(self):
with pytest.raises(UrlSecurityError, match="端口"):
validate_url_safety("http://example.com:3306/file")
class TestSafeDownload(unittest.TestCase):
"""安全下载函数测试."""
# ── 内网主机名 / SSRF 防护 ───────────────────────────────────────────────────
def setUp(self):
self.temp_dir = tempfile.mkdtemp()
def tearDown(self):
shutil.rmtree(self.temp_dir, ignore_errors=True)
class TestInternalHostnameProtection:
"""内网主机名防护测试"""
def test_safe_download_file_rejects_ssrf(self):
"""SSRF 风险 URL 应该被拒绝下载."""
dest = os.path.join(self.temp_dir, "test.bin")
with self.assertRaises(UrlSecurityError):
safe_download_file("http://127.0.0.1/test", dest, purpose="test")
def test_localhost_raises(self):
with pytest.raises(UrlSecurityError, match="内部主机名"):
validate_url_safety("http://localhost/file")
def test_safe_download_bytes_rejects_ssrf(self):
"""SSRF 风险 URL 应该被拒绝下载(bytes 版本)."""
with self.assertRaises(UrlSecurityError):
safe_download_bytes("http://localhost/test", purpose="test")
def test_localhost_mixed_case_raises(self):
with pytest.raises(UrlSecurityError):
validate_url_safety("http://LocalHost/file")
def test_safe_download_file_size_limit(self):
"""超过大小限制应该被拒绝."""
# 用 mock server 测试太大的 content-length
dest = os.path.join(self.temp_dir, "test.bin")
# 直接验证参数:max_size=0 时任何下载都应超限
# (这里用一个可访问的 URL 并设置极小的限制)
# 为避免依赖外部网络,这里只测试函数参数传递
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
mock_resp = unittest.mock.MagicMock()
mock_resp.headers = {"Content-Length": "1000"}
mock_resp.read.return_value = b""
mock_opener.return_value.open.return_value = mock_resp
# 设置 max_size=500content-length=1000 应被拒绝
with self.assertRaises(UrlSecurityError):
safe_download_file(
"https://example.com/test",
def test_localhost_localdomain_raises(self):
with pytest.raises(UrlSecurityError):
_check_internal_hostnames("localhost.localdomain")
def test_metadata_hostname_raises(self):
with pytest.raises(UrlSecurityError):
_check_internal_hostnames("metadata")
def test_metadata_google_internal_raises(self):
with pytest.raises(UrlSecurityError):
_check_internal_hostnames("metadata.google.internal")
def test_dot_local_domain_raises(self):
with pytest.raises(UrlSecurityError, match="内网域名"):
validate_url_safety("http://myservice.local/file")
def test_dot_internal_domain_raises(self):
with pytest.raises(UrlSecurityError, match="内网域名"):
validate_url_safety("http://myservice.internal/file")
def test_dot_localdomain_raises(self):
with pytest.raises(UrlSecurityError):
_check_internal_hostnames("server.localdomain")
def test_loopback_ip_127_0_0_1_raises(self):
with pytest.raises(UrlSecurityError, match="直接 IP|回环"):
validate_url_safety("http://127.0.0.1/file")
def test_metadata_ip_169_254_raises(self):
"""云元数据服务 IP"""
with pytest.raises(UrlSecurityError):
validate_url_safety("http://169.254.169.254/latest/meta-data/")
def test_private_ip_10_raises(self):
with pytest.raises(UrlSecurityError):
validate_url_safety("http://10.0.0.1/file")
def test_private_ip_172_16_raises(self):
with pytest.raises(UrlSecurityError):
validate_url_safety("http://172.16.0.1/file")
def test_private_ip_192_168_raises(self):
with pytest.raises(UrlSecurityError):
validate_url_safety("http://192.168.1.1/file")
def test_unspecified_ip_0_0_0_0_raises(self):
with pytest.raises(UrlSecurityError):
validate_url_safety("http://0.0.0.0/file")
def test_ipv6_loopback_raises(self):
with pytest.raises(UrlSecurityError):
validate_url_safety("http://[::1]/file")
def test_public_ip_ok(self):
"""公网IP在ALLOW_DIRECT_IP默认关闭时应被拦截"""
# 默认 ALLOW_DIRECT_IP = false
with pytest.raises(UrlSecurityError, match="直接 IP"):
validate_url_safety("http://8.8.8.8/file")
# ── 可信域名白名单 ───────────────────────────────────────────────────────────
class TestTrustedDomains:
"""可信域名白名单测试"""
def test_is_trusted_domain_exact_match(self):
with patch("packages.shared.url_security.TRUSTED_DOMAINS", {"example.com", "cdn.example.org"}):
# 重新加载模块以应用环境变量不太现实,直接测函数
# 直接改全局状态再还原
import packages.shared.url_security as mod
from packages.shared.url_security import _is_trusted_domain
original = mod.TRUSTED_DOMAINS
mod.TRUSTED_DOMAINS = {"example.com", "cdn.example.org"}
try:
assert _is_trusted_domain("example.com") is True
assert _is_trusted_domain("cdn.example.org") is True
finally:
mod.TRUSTED_DOMAINS = original
def test_is_trusted_domain_subdomain(self):
import packages.shared.url_security as mod
original = mod.TRUSTED_DOMAINS
mod.TRUSTED_DOMAINS = {"example.com"}
try:
assert mod._is_trusted_domain("sub.example.com") is True
assert mod._is_trusted_domain("a.b.example.com") is True
finally:
mod.TRUSTED_DOMAINS = original
def test_is_trusted_domain_no_match(self):
import packages.shared.url_security as mod
original = mod.TRUSTED_DOMAINS
mod.TRUSTED_DOMAINS = {"example.com"}
try:
assert mod._is_trusted_domain("other.com") is False
assert mod._is_trusted_domain("notexample.com") is False
finally:
mod.TRUSTED_DOMAINS = original
def test_validate_with_trusted_domains_restricted(self):
"""白名单非空时,不在白名单中的域名被拒"""
import packages.shared.url_security as mod
original = mod.TRUSTED_DOMAINS
mod.TRUSTED_DOMAINS = {"trusted.com"}
try:
# 不在白名单中 - 在 _is_trusted_domain 检查时就被拒,不走 DNS
with pytest.raises(UrlSecurityError, match="白名单"):
validate_url_safety("https://untrusted.com/file")
# 在白名单中 - 需要 mock DNS 解析避免实际网络请求
with patch("packages.shared.url_security._check_ssrf_domain"):
result = validate_url_safety("https://trusted.com/file")
assert result == "https://trusted.com/file"
# 子域名
result = validate_url_safety("https://sub.trusted.com/file")
assert result == "https://sub.trusted.com/file"
finally:
mod.TRUSTED_DOMAINS = original
# ── is_url_safe 便捷函数 ─────────────────────────────────────────────────────
class TestIsUrlSafe:
"""is_url_safe 便捷函数测试"""
def test_safe_url_returns_true(self):
assert is_url_safe("https://example.com/file") is True
def test_unsafe_url_returns_false(self):
assert is_url_safe("http://localhost/file") is False
def test_empty_url_returns_false(self):
assert is_url_safe("") is False
def test_invalid_scheme_returns_false(self):
assert is_url_safe("ftp://example.com/file") is False
# ── UrlSecurityError 异常类 ───────────────────────────────────────────────────
class TestUrlSecurityError:
"""UrlSecurityError 异常类测试"""
def test_is_value_error_subclass(self):
assert issubclass(UrlSecurityError, ValueError)
def test_error_message(self):
err = UrlSecurityError("test message")
assert str(err) == "test message"
# ── NoRedirectHandler ────────────────────────────────────────────────────────
class TestNoRedirectHandler:
"""NoRedirectHandler 测试"""
def test_redirect_request_returns_none(self):
handler = NoRedirectHandler()
result = handler.redirect_request(
MagicMock(),
MagicMock(),
302,
"Found",
{"Location": "http://other.com"},
"http://other.com",
)
assert result is None
# ── 魔数校验 ─────────────────────────────────────────────────────────────────
class TestMagicNumberValidation:
"""文件魔数校验测试"""
def test_valid_png(self, tmp_path):
f = tmp_path / "test.png"
# PNG 文件头: 89 50 4E 47 0D 0A 1A 0A
f.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
# 不抛异常 = 通过
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_valid_jpeg(self, tmp_path):
f = tmp_path / "test.jpg"
# JPEG 文件头: FF D8 FF
f.write_bytes(b"\xff\xd8\xff\xe0" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_valid_gif87a(self, tmp_path):
f = tmp_path / "test.gif"
f.write_bytes(b"GIF87a" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_valid_gif89a(self, tmp_path):
f = tmp_path / "test.gif"
f.write_bytes(b"GIF89a" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_valid_webp(self, tmp_path):
f = tmp_path / "test.webp"
# RIFF....WEBP
data = bytearray(b"RIFF")
data += b"\x00\x00\x00\x00" # size placeholder
data += b"WEBP"
data += b"\x00" * 100
f.write_bytes(bytes(data))
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_valid_bmp(self, tmp_path):
f = tmp_path / "test.bmp"
f.write_bytes(b"BM" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_valid_wav(self, tmp_path):
f = tmp_path / "test.wav"
# RIFF....WAVE
data = bytearray(b"RIFF")
data += b"\x00\x00\x00\x00"
data += b"WAVE"
data += b"\x00" * 100
f.write_bytes(bytes(data))
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
def test_valid_mp3_id3(self, tmp_path):
f = tmp_path / "test.mp3"
f.write_bytes(b"ID3\x03\x00\x00\x00\x00\x00\x00" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
def test_valid_mp3_adts(self, tmp_path):
f = tmp_path / "test.mp3"
f.write_bytes(b"\xff\xfb\x90\x00" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
def test_valid_ogg(self, tmp_path):
f = tmp_path / "test.ogg"
f.write_bytes(b"OggS\x00\x02\x00\x00" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
def test_valid_flac(self, tmp_path):
f = tmp_path / "test.flac"
f.write_bytes(b"fLaC" + b"\x00" * 100)
_validate_magic_number(str(f), ALLOWED_AUDIO_MIME_TYPES)
def test_invalid_file_content_raises(self, tmp_path):
f = tmp_path / "test.bin"
f.write_bytes(b"this is not an image file at all")
with pytest.raises(UrlSecurityError, match="魔数"):
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_empty_file_raises(self, tmp_path):
f = tmp_path / "empty.bin"
f.write_bytes(b"")
with pytest.raises(UrlSecurityError, match="为空"):
_validate_magic_number(str(f), ALLOWED_IMAGE_MIME_TYPES)
def test_nonexistent_file_raises(self, tmp_path):
with pytest.raises(UrlSecurityError, match="读取文件头失败"):
_validate_magic_number(str(tmp_path / "no_such_file"), ALLOWED_IMAGE_MIME_TYPES)
def test_no_allowed_mime_types_skips(self, tmp_path):
"""allowed_mime_types 为空时跳过校验"""
f = tmp_path / "test.bin"
f.write_bytes(b"random data here")
# 不抛异常
_validate_magic_number(str(f), set())
def test_unknown_mime_types_skips(self, tmp_path):
"""没有已知魔数的 MIME 类型跳过校验"""
f = tmp_path / "test.bin"
f.write_bytes(b"random data")
_validate_magic_number(str(f), {"application/x-unknown-type"})
# ── safe_download_file (mock 网络) ───────────────────────────────────────────
class TestSafeDownloadFile:
"""safe_download_file 下载测试(mock 网络)"""
def test_download_success(self, tmp_path):
test_content = b"Hello, this is test file content!"
dest = str(tmp_path / "output.bin")
mock_resp = MagicMock()
mock_resp.headers = {"Content-Type": "application/octet-stream"}
mock_resp.read.side_effect = [test_content, b""]
with patch("packages.shared.url_security.NoRedirectHandler") as mock_handler_cls:
mock_handler = MagicMock()
mock_handler_cls.return_value = mock_handler
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
with patch("urllib.request.build_opener", return_value=mock_opener):
size = safe_download_file(
"https://example.com/test.bin",
dest,
purpose="test",
max_size=500,
)
def test_safe_download_file_mime_rejected(self):
"""不允许的 MIME 类型应该被拒绝."""
dest = os.path.join(self.temp_dir, "test.bin")
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
mock_resp = unittest.mock.MagicMock()
mock_resp.headers = {"Content-Type": "text/html"}
mock_resp.read.return_value = b""
mock_opener.return_value.open.return_value = mock_resp
with self.assertRaises(UrlSecurityError):
safe_download_file(
"https://example.com/test.mp3",
dest,
purpose="test",
allowed_mime_types=ALLOWED_AUDIO_MIME_TYPES,
)
assert size == len(test_content)
with open(dest, "rb") as f:
assert f.read() == test_content
def test_download_with_mime_check_passes(self, tmp_path):
# PNG 文件
test_content = b"\x89PNG\r\n\x1a\n" + b"\x00" * 200
dest = str(tmp_path / "test.png")
mock_resp = MagicMock()
mock_resp.headers = {"Content-Type": "image/png"}
mock_resp.read.side_effect = [test_content, b""]
with patch("urllib.request.build_opener") as mock_build:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_build.return_value = mock_opener
def test_safe_download_file_mime_allowed(self):
"""允许的 MIME 类型应该通过."""
dest = os.path.join(self.temp_dir, "test.mp3")
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
mock_resp = unittest.mock.MagicMock()
mock_resp.headers = {"Content-Type": "audio/mpeg"}
mock_resp.read.side_effect = [b"ID3audio_data", b""]
mock_resp.geturl.return_value = "https://example.com/test.mp3"
mock_opener.return_value.open.return_value = mock_resp
size = safe_download_file(
"https://example.com/test.mp3",
"https://example.com/test.png",
dest,
purpose="test",
allowed_mime_types=ALLOWED_AUDIO_MIME_TYPES,
allowed_mime_types={"image/png", "image/jpeg"},
)
self.assertEqual(size, 13)
self.assertTrue(os.path.exists(dest))
def test_safe_download_file_stream_size_limit(self):
"""流式下载时超过大小限制应该中断."""
dest = os.path.join(self.temp_dir, "test.bin")
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
mock_resp = unittest.mock.MagicMock()
mock_resp.headers = {}
# 每次返回 100 字节,max_size=500,第 6 次读取就超限
mock_resp.read.side_effect = lambda n: b"x" * n if n < 1000 else b"x" * 100
# 改成返回固定 100 字节,直到第 N 次后返回空
call_count = [0]
assert size == len(test_content)
def mock_read(size):
call_count[0] += 1
if call_count[0] > 10:
return b""
return b"x" * 100
def test_download_mime_type_rejected(self, tmp_path):
test_content = b"GIF89a" + b"\x00" * 50
dest = str(tmp_path / "test.gif")
mock_resp.read = mock_read
mock_opener.return_value.open.return_value = mock_resp
with self.assertRaises(UrlSecurityError):
mock_resp = MagicMock()
mock_resp.headers = {"Content-Type": "image/gif"}
mock_resp.read.side_effect = [test_content, b""]
with patch("urllib.request.build_opener") as mock_build:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_build.return_value = mock_opener
with pytest.raises(UrlSecurityError, match="Content-Type"):
safe_download_file(
"https://example.com/test",
"https://example.com/test.gif",
dest,
purpose="test",
max_size=500, # 500 字节上限
allowed_mime_types={"image/png"},
)
def test_safe_download_bytes_returns_content(self):
"""safe_download_bytes 应该返回文件内容."""
test_data = b"ID3hello world test audio"
with unittest.mock.patch("urllib.request.build_opener") as mock_opener:
mock_resp = unittest.mock.MagicMock()
mock_resp.headers = {"Content-Type": "audio/mpeg"}
call_count = [0]
def test_download_size_limit_exceeded(self, tmp_path):
"""流式下载时超过大小限制被中断(无 Content-Length header"""
dest = str(tmp_path / "big.bin")
chunk = b"x" * 1024 # 1KB chunks
def mock_read(size):
call_count[0] += 1
if call_count[0] > 1:
return b""
return test_data
mock_resp = MagicMock()
# 没有 Content-Length header,走流式检查
mock_resp.headers = {"Content-Type": "application/octet-stream"}
# 模拟多次读取,超过 5KB 限制(6个chunk = 6KB
mock_resp.read.side_effect = [chunk] * 6 + [b""]
with patch("urllib.request.build_opener") as mock_build:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_build.return_value = mock_opener
with pytest.raises(UrlSecurityError, match="超过大小限制"):
safe_download_file(
"https://example.com/big.bin",
dest,
purpose="test",
max_size=5000, # 5KB limit
)
def test_download_content_length_too_large(self, tmp_path):
dest = str(tmp_path / "big.bin")
mock_resp = MagicMock()
mock_resp.headers = {"Content-Type": "application/octet-stream", "Content-Length": "1000000"}
mock_resp.read.side_effect = [b"data"]
with patch("urllib.request.build_opener") as mock_build:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_build.return_value = mock_opener
with pytest.raises(UrlSecurityError, match="文件过大"):
safe_download_file(
"https://example.com/big.bin",
dest,
purpose="test",
max_size=500000,
)
def test_download_localhost_rejected(self, tmp_path):
"""内网 URL 在下载前就被拒"""
dest = str(tmp_path / "out.bin")
with pytest.raises(UrlSecurityError):
safe_download_file("http://localhost/file", dest)
def test_download_invalid_scheme_rejected(self, tmp_path):
dest = str(tmp_path / "out.bin")
with pytest.raises(UrlSecurityError, match="scheme"):
safe_download_file("ftp://example.com/file", dest)
# ── safe_download_bytes ───────────────────────────────────────────────────────
class TestSafeDownloadBytes:
"""safe_download_bytes 测试"""
def test_download_returns_bytes(self, tmp_path):
test_content = b"hello bytes download test"
mock_resp = MagicMock()
mock_resp.headers = {"Content-Type": "application/octet-stream"}
mock_resp.read.side_effect = [test_content, b""]
with patch("urllib.request.build_opener") as mock_build:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_build.return_value = mock_opener
mock_resp.read = mock_read
mock_opener.return_value.open.return_value = mock_resp
result = safe_download_bytes(
"https://example.com/test.mp3",
"https://example.com/test.bin",
purpose="test",
allowed_mime_types=ALLOWED_AUDIO_MIME_TYPES,
)
self.assertEqual(result, test_data)
assert result == test_content
assert isinstance(result, bytes)
def test_download_unsafe_url_raises(self):
with pytest.raises(UrlSecurityError):
safe_download_bytes("http://127.0.0.1/secret")
# ── 常量导出验证 ─────────────────────────────────────────────────────────────
class TestConstants:
"""模块常量验证"""
def test_allowed_schemes(self):
assert "http" in ALLOWED_SCHEMES
assert "https" in ALLOWED_SCHEMES
assert len(ALLOWED_SCHEMES) == 2
def test_allowed_ports(self):
assert 80 in ALLOWED_PORTS
assert 443 in ALLOWED_PORTS
assert len(ALLOWED_PORTS) == 2
def test_max_url_length(self):
assert MAX_URL_LENGTH == 2048