Files
model-platform/backend/tests/test_resources.py
T
2026-08-14 20:37:25 +08:00

346 lines
12 KiB
Python

"""Unit tests for data-resource path derivation and helpers.
These tests do not need a database because they exercise pure helpers.
"""
from __future__ import annotations
import datetime
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from backend.resources import (
compute_jupyter_relative_path,
resource_directory,
resource_payload,
)
from backend.services.storage import _safe_file_name, _safe_path_segment
_BIND_WS = "01WS0000000000000000000A"
_BIND_USER = "01USR0000000000000000000A"
class _ExecuteResult:
def __init__(self, rows):
self._rows = rows
def all(self):
return self._rows
class _BindSessionMock:
"""Mocked session for bind_resource.
``new_object_key`` is the StorageObjects row of the upload being bound;
``same_name_rows`` are (DataResources, StorageObjects) candidates already
in the workspace with the same resource_name.
"""
def __init__(self, new_object_key: str, same_name_rows=(), reused_resource=None):
self._new_object_key = new_object_key
self._same_name_rows = list(same_name_rows)
self._reused_resource = reused_resource
self._scalar_calls = 0
async def scalar(self, _stmt):
self._scalar_calls += 1
# 1st scalar: UploadSessions lookup; 2nd: storage_object_id reuse check.
if self._scalar_calls == 1:
return SimpleNamespace(
upload_id="01UPL0000000000000000000B",
storage_object_id="01OBJ0000000000000000000B",
workspace_id=_BIND_WS,
user_id=_BIND_USER,
)
if self._scalar_calls == 2 and self._reused_resource is not None:
return self._reused_resource
return None
async def get(self, _model, _pk):
return _make_storage_object(self._new_object_key)
async def execute(self, _stmt):
# 模拟 SQL 的 owner 过滤与 self-exclusion:只返回属于当前用户且
# 不是当前 upload 已绑定对象的同名候选。
rows = [
(resource, storage_object)
for resource, storage_object in self._same_name_rows
if resource.owner_user_id == _BIND_USER
and resource.storage_object_id != "01OBJ0000000000000000000B"
]
return _ExecuteResult(rows)
def add(self, _obj):
pass
async def flush(self):
pass
async def refresh(self, _obj, attribute_names=None):
_obj.created_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
_obj.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
def _bind_context():
return SimpleNamespace(
request_id="01REQ0000000000000000000A",
user=SimpleNamespace(user_id=_BIND_USER),
workspace=SimpleNamespace(workspace_id=_BIND_WS),
)
def _bind_payload():
return SimpleNamespace(
resource_name="data.csv", description=None, visibility="private"
)
@pytest.mark.asyncio
async def test_bind_resource_rejects_duplicate_name_in_same_directory() -> None:
"""Same resource_name in the same directory raises 409."""
from backend.resources import bind_resource
existing_rows = [
(
_make_resource(_BIND_WS, _BIND_USER),
_make_storage_object(f"{_BIND_WS}/{_BIND_USER}/data.csv"),
)
]
session = _BindSessionMock(
new_object_key=f"{_BIND_WS}/{_BIND_USER}/data.csv",
same_name_rows=existing_rows,
)
with pytest.raises(Exception) as exc_info:
await bind_resource(
upload_id="01UPL0000000000000000000B",
payload=_bind_payload(),
request=MagicMock(),
context=_bind_context(),
session=session,
)
assert exc_info.value.status_code == 409
assert "already exists" in exc_info.value.detail
@pytest.mark.asyncio
async def test_bind_resource_allows_same_name_in_different_directory() -> None:
"""Same resource_name in a different directory binds successfully."""
from backend.resources import bind_resource
existing_rows = [
(
_make_resource(_BIND_WS, _BIND_USER),
_make_storage_object(f"{_BIND_WS}/{_BIND_USER}/data.csv"),
)
]
session = _BindSessionMock(
new_object_key=f"{_BIND_WS}/{_BIND_USER}/subdir/data.csv",
same_name_rows=existing_rows,
)
result = await bind_resource(
upload_id="01UPL0000000000000000000B",
payload=_bind_payload(),
request=MagicMock(),
context=_bind_context(),
session=session,
)
assert result["data"]["resource_name"] == "data.csv"
assert result["data"]["jupyter_accessible_path"] == "subdir/data.csv"
@pytest.mark.asyncio
async def test_bind_resource_allows_same_name_when_workspace_empty() -> None:
"""No same-name rows at all: bind succeeds (root directory)."""
from backend.resources import bind_resource
session = _BindSessionMock(
new_object_key=f"{_BIND_WS}/{_BIND_USER}/data.csv",
)
result = await bind_resource(
upload_id="01UPL0000000000000000000A",
payload=_bind_payload(),
request=MagicMock(),
context=_bind_context(),
session=session,
)
assert result["data"]["resource_name"] == "data.csv"
@pytest.mark.asyncio
async def test_bind_resource_allows_same_name_for_different_owner() -> None:
"""其他用户在同目录下的同名资源不阻塞当前用户的绑定。"""
from backend.resources import bind_resource
other_user = "01USR0000000000000000000B"
existing_rows = [
(
_make_resource(_BIND_WS, other_user),
_make_storage_object(f"{_BIND_WS}/{other_user}/data.csv"),
)
]
session = _BindSessionMock(
new_object_key=f"{_BIND_WS}/{_BIND_USER}/data.csv",
same_name_rows=existing_rows,
)
result = await bind_resource(
upload_id="01UPL0000000000000000000C",
payload=_bind_payload(),
request=MagicMock(),
context=_bind_context(),
session=session,
)
assert result["data"]["resource_name"] == "data.csv"
@pytest.mark.asyncio
async def test_bind_resource_allows_rebinding_same_storage_object() -> None:
"""重新绑定同一 upload_id 应走 idempotent 复用路径,不触发 409。"""
from backend.resources import bind_resource
new_object_key = f"{_BIND_WS}/{_BIND_USER}/data.csv"
existing_resource = _make_resource(_BIND_WS, _BIND_USER)
existing_resource.resource_name = "data.csv"
existing_resource.storage_object_id = "01OBJ0000000000000000000B"
existing_object = _make_storage_object(new_object_key)
existing_object.storage_object_id = "01OBJ0000000000000000000B"
same_name_rows = [(existing_resource, existing_object)]
session = _BindSessionMock(
new_object_key=new_object_key,
same_name_rows=same_name_rows,
reused_resource=existing_resource,
)
result = await bind_resource(
upload_id="01UPL0000000000000000000B",
payload=_bind_payload(),
request=MagicMock(),
context=_bind_context(),
session=session,
)
assert result["meta"]["reused"] is True
assert result["data"]["resource_name"] == "data.csv"
def test_safe_path_segment_cleans_special_characters():
assert _safe_path_segment("train") == "train"
assert _safe_path_segment("train v1") == "train_v1"
assert _safe_path_segment("../foo") == "foo"
assert _safe_path_segment("a/b") == "a_b"
assert _safe_path_segment("...") == "untitled"
def test_safe_file_name_rejects_traversal_and_hidden():
assert _safe_file_name("data.csv") == "data.csv"
with pytest.raises(Exception):
_safe_file_name("../data.csv")
with pytest.raises(Exception):
_safe_file_name(".hidden.csv")
def _make_resource(workspace_id: str, owner_user_id: str):
return SimpleNamespace(
resource_id="01RES0000000000000000000A",
workspace_id=workspace_id,
storage_object_id="01OBJ0000000000000000000A",
owner_user_id=owner_user_id,
resource_name="sample",
description=None,
visibility="workspace",
status="active",
created_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
updated_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
)
def _make_storage_object(object_key: str):
return SimpleNamespace(
storage_object_id="01OBJ0000000000000000000A",
object_key=object_key,
file_name="data.csv",
file_extension=".csv",
mime_type="text/csv",
size_bytes=42,
content_hash="a" * 64,
object_status="available",
)
def test_resource_payload_legacy_dot_resources():
ws = "01WS00000000000000000000A"
user = "01USR000000000000000000A"
payload = resource_payload(
_make_resource(ws, user),
_make_storage_object(f"{ws}/{user}/.resources/data.csv"),
)
assert payload["jupyter_accessible_path"] == ".resources/data.csv"
assert payload["absolute_path"].endswith(f"{ws}/{user}/.resources/data.csv")
def test_resource_payload_new_flat_path():
ws = "01WS00000000000000000000A"
user = "01USR000000000000000000A"
payload = resource_payload(
_make_resource(ws, user),
_make_storage_object(f"{ws}/{user}/data.csv"),
)
assert payload["jupyter_accessible_path"] == "data.csv"
assert payload["absolute_path"].endswith(f"{ws}/{user}/data.csv")
def test_resource_payload_new_nested_path():
ws = "01WS00000000000000000000A"
user = "01USR000000000000000000A"
payload = resource_payload(
_make_resource(ws, user),
_make_storage_object(f"{ws}/{user}/train/v1/data.csv"),
)
assert payload["jupyter_accessible_path"] == "train/v1/data.csv"
assert payload["absolute_path"].endswith(f"{ws}/{user}/train/v1/data.csv")
def test_compute_jupyter_relative_path_for_legacy_and_new_paths():
# Legacy .resources path still resolves correctly.
assert compute_jupyter_relative_path("notebooks/exp.ipynb", ".resources/data.csv") == "../.resources/data.csv"
# New nested path resolves relative to the script directory.
assert compute_jupyter_relative_path("notebooks/exp.ipynb", "train/v1/data.csv") == "../train/v1/data.csv"
# Same-directory script.
assert compute_jupyter_relative_path("exp.ipynb", "data.csv") == "data.csv"
def test_resource_directory_parses_object_key():
ws = "01WS00000000000000000000A"
user = "01USR000000000000000000A"
# 根目录文件:目录为 ""。
assert resource_directory(f"{ws}/{user}/data.csv", ws, user) == ""
# 子目录 / 多级目录。
assert resource_directory(f"{ws}/{user}/sub/data.csv", ws, user) == "sub"
assert (
resource_directory(f"{ws}/{user}/train/v1/data.csv", ws, user)
== "train/v1"
)
# 旧版 .resources 布局:目录为 ".resources"。
assert (
resource_directory(f"{ws}/{user}/.resources/data.csv", ws, user)
== ".resources"
)
# 不匹配 ws/user 前缀的键按根目录处理。
assert resource_directory("other-bucket-key.csv", ws, user) == ""
def test_data_resources_model_allows_duplicate_storage_object_reference() -> None:
"""With uk_data_resources_object dropped, no unique index covers
storage_object_id, so multiple DataResources rows may reference the
same storage object.
"""
from common.db.models import DataResources
index_names = {idx.name for idx in DataResources.__table__.indexes}
assert "uk_data_resources_object" not in index_names
for idx in DataResources.__table__.indexes:
if idx.unique:
cols = {c.name for c in idx.columns}
assert "storage_object_id" not in cols, (
f"unexpected unique index {idx.name} on storage_object_id"
)