mirror of
https://github.com/langgenius/dify.git
synced 2026-03-27 01:00:13 +08:00
236 lines
8.1 KiB
Python
236 lines
8.1 KiB
Python
import importlib
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from core.workflow.file_reference import build_file_reference
|
|
from graphon.file import FILE_MODEL_IDENTITY, FileTransferMethod
|
|
from models.model import Conversation, Message
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def patch_file_helpers(monkeypatch: pytest.MonkeyPatch):
|
|
"""
|
|
Patch file_helpers.get_signed_file_url to a deterministic stub.
|
|
"""
|
|
model_module = importlib.import_module("models.model")
|
|
dummy = types.SimpleNamespace(get_signed_file_url=lambda fid: f"https://signed.example/{fid}")
|
|
# Inject/override file_helpers on models.model
|
|
monkeypatch.setattr(model_module, "file_helpers", dummy, raising=False)
|
|
|
|
|
|
def _wrap_md(url: str) -> str:
|
|
"""
|
|
Wrap a raw URL into the markdown that re_sign_file_url_answer expects:
|
|
[link](<url>)
|
|
"""
|
|
return f"please click [file]({url}) to download."
|
|
|
|
|
|
def test_file_preview_valid_replaced():
|
|
"""
|
|
Valid file-preview URL must be re-signed:
|
|
- Extract upload_file_id correctly
|
|
- Replace the original URL with the signed URL
|
|
"""
|
|
upload_id = "abc-123"
|
|
url = f"/files/{upload_id}/file-preview?timestamp=111&nonce=222&sign=333"
|
|
msg = Message(answer=_wrap_md(url))
|
|
|
|
out = msg.re_sign_file_url_answer
|
|
assert f"https://signed.example/{upload_id}" in out
|
|
assert url not in out
|
|
|
|
|
|
def test_file_preview_misspelled_not_replaced():
|
|
"""
|
|
Misspelled endpoint 'file-previe?timestamp=' should NOT be rewritten.
|
|
"""
|
|
upload_id = "zzz-001"
|
|
# path deliberately misspelled: file-previe? (missing 'w')
|
|
# and we append ¬e=file-preview to trick the old `"file-preview" in url` check.
|
|
url = f"/files/{upload_id}/file-previe?timestamp=111&nonce=222&sign=333¬e=file-preview"
|
|
original = _wrap_md(url)
|
|
msg = Message(answer=original)
|
|
|
|
out = msg.re_sign_file_url_answer
|
|
# Expect NO replacement, should not rewrite misspelled file-previe URL
|
|
assert out == original
|
|
|
|
|
|
def test_image_preview_valid_replaced():
|
|
"""
|
|
Valid image-preview URL must be re-signed.
|
|
"""
|
|
upload_id = "img-789"
|
|
url = f"/files/{upload_id}/image-preview?timestamp=123&nonce=456&sign=789"
|
|
msg = Message(answer=_wrap_md(url))
|
|
|
|
out = msg.re_sign_file_url_answer
|
|
assert f"https://signed.example/{upload_id}" in out
|
|
assert url not in out
|
|
|
|
|
|
def test_image_preview_misspelled_not_replaced():
|
|
"""
|
|
Misspelled endpoint 'image-previe?timestamp=' should NOT be rewritten.
|
|
"""
|
|
upload_id = "img-err-42"
|
|
url = f"/files/{upload_id}/image-previe?timestamp=1&nonce=2&sign=3¬e=image-preview"
|
|
original = _wrap_md(url)
|
|
msg = Message(answer=original)
|
|
|
|
out = msg.re_sign_file_url_answer
|
|
# Expect NO replacement, should not rewrite misspelled image-previe URL
|
|
assert out == original
|
|
|
|
|
|
def _build_local_file_mapping(record_id: str, *, tenant_id: str | None = None) -> dict[str, object]:
|
|
mapping: dict[str, object] = {
|
|
"dify_model_identity": FILE_MODEL_IDENTITY,
|
|
"transfer_method": FileTransferMethod.LOCAL_FILE,
|
|
"reference": build_file_reference(record_id=record_id),
|
|
"type": "document",
|
|
"filename": "example.txt",
|
|
"extension": ".txt",
|
|
"mime_type": "text/plain",
|
|
"size": 1,
|
|
}
|
|
if tenant_id is not None:
|
|
mapping["tenant_id"] = tenant_id
|
|
return mapping
|
|
|
|
|
|
@pytest.mark.parametrize("owner_cls", [Conversation, Message])
|
|
def test_inputs_resolve_owner_tenant_for_single_file_mapping(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
owner_cls: type[Conversation] | type[Message],
|
|
):
|
|
model_module = importlib.import_module("models.model")
|
|
build_calls: list[tuple[dict[str, object], str]] = []
|
|
|
|
monkeypatch.setattr(model_module.db.session, "scalar", lambda _: "tenant-from-app")
|
|
|
|
def fake_build_from_mapping(*, mapping, tenant_id, config=None, strict_type_validation=False, access_controller):
|
|
_ = config, strict_type_validation, access_controller
|
|
build_calls.append((dict(mapping), tenant_id))
|
|
return {"tenant_id": tenant_id, "upload_file_id": mapping.get("upload_file_id")}
|
|
|
|
monkeypatch.setattr("factories.file_factory.build_from_mapping", fake_build_from_mapping)
|
|
|
|
owner = owner_cls(app_id="app-1")
|
|
owner.inputs = {"file": _build_local_file_mapping("upload-1")}
|
|
|
|
restored_inputs = owner.inputs
|
|
|
|
assert restored_inputs["file"] == {"tenant_id": "tenant-from-app", "upload_file_id": "upload-1"}
|
|
assert build_calls == [
|
|
(
|
|
{
|
|
**_build_local_file_mapping("upload-1"),
|
|
"upload_file_id": "upload-1",
|
|
},
|
|
"tenant-from-app",
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("owner_cls", [Conversation, Message])
|
|
def test_inputs_resolve_owner_tenant_for_file_list_mapping(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
owner_cls: type[Conversation] | type[Message],
|
|
):
|
|
model_module = importlib.import_module("models.model")
|
|
build_calls: list[tuple[dict[str, object], str]] = []
|
|
|
|
monkeypatch.setattr(model_module.db.session, "scalar", lambda _: "tenant-from-app")
|
|
|
|
def fake_build_from_mapping(*, mapping, tenant_id, config=None, strict_type_validation=False, access_controller):
|
|
_ = config, strict_type_validation, access_controller
|
|
build_calls.append((dict(mapping), tenant_id))
|
|
return {"tenant_id": tenant_id, "upload_file_id": mapping.get("upload_file_id")}
|
|
|
|
monkeypatch.setattr("factories.file_factory.build_from_mapping", fake_build_from_mapping)
|
|
|
|
owner = owner_cls(app_id="app-1")
|
|
owner.inputs = {
|
|
"files": [
|
|
_build_local_file_mapping("upload-1"),
|
|
_build_local_file_mapping("upload-2"),
|
|
]
|
|
}
|
|
|
|
restored_inputs = owner.inputs
|
|
|
|
assert restored_inputs["files"] == [
|
|
{"tenant_id": "tenant-from-app", "upload_file_id": "upload-1"},
|
|
{"tenant_id": "tenant-from-app", "upload_file_id": "upload-2"},
|
|
]
|
|
assert build_calls == [
|
|
(
|
|
{
|
|
**_build_local_file_mapping("upload-1"),
|
|
"upload_file_id": "upload-1",
|
|
},
|
|
"tenant-from-app",
|
|
),
|
|
(
|
|
{
|
|
**_build_local_file_mapping("upload-2"),
|
|
"upload_file_id": "upload-2",
|
|
},
|
|
"tenant-from-app",
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("owner_cls", [Conversation, Message])
|
|
def test_inputs_prefer_serialized_tenant_id_when_present(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
owner_cls: type[Conversation] | type[Message],
|
|
):
|
|
model_module = importlib.import_module("models.model")
|
|
|
|
def fail_if_called(_):
|
|
raise AssertionError("App tenant lookup should not run when tenant_id exists in the file mapping")
|
|
|
|
monkeypatch.setattr(model_module.db.session, "scalar", fail_if_called)
|
|
|
|
def fake_build_from_mapping(*, mapping, tenant_id, config=None, strict_type_validation=False, access_controller):
|
|
_ = config, strict_type_validation, access_controller
|
|
return {"tenant_id": tenant_id, "upload_file_id": mapping.get("upload_file_id")}
|
|
|
|
monkeypatch.setattr("factories.file_factory.build_from_mapping", fake_build_from_mapping)
|
|
|
|
owner = owner_cls(app_id="app-1")
|
|
owner.inputs = {"file": _build_local_file_mapping("upload-1", tenant_id="tenant-from-payload")}
|
|
|
|
restored_inputs = owner.inputs
|
|
|
|
assert restored_inputs["file"] == {
|
|
"tenant_id": "tenant-from-payload",
|
|
"upload_file_id": "upload-1",
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("owner_cls", [Conversation, Message])
|
|
def test_inputs_restore_external_remote_url_file_mappings(owner_cls: type[Conversation] | type[Message]) -> None:
|
|
owner = owner_cls(app_id="app-1")
|
|
owner.inputs = {
|
|
"file": {
|
|
"dify_model_identity": FILE_MODEL_IDENTITY,
|
|
"transfer_method": FileTransferMethod.REMOTE_URL,
|
|
"type": "document",
|
|
"url": "https://example.com/report.pdf",
|
|
"filename": "report.pdf",
|
|
"extension": ".pdf",
|
|
"mime_type": "application/pdf",
|
|
"size": 1,
|
|
}
|
|
}
|
|
|
|
restored_file = owner.inputs["file"]
|
|
|
|
assert restored_file.transfer_method == FileTransferMethod.REMOTE_URL
|
|
assert restored_file.remote_url == "https://example.com/report.pdf"
|