mirror of
https://github.com/langgenius/dify.git
synced 2026-11-10 13:26:58 +08:00
310 lines
9.4 KiB
Python
310 lines
9.4 KiB
Python
"""Unit tests for the @accepts / @returns contract decorators.
|
|
|
|
Exercises the decorators in isolation (not through a real controller): a plain
|
|
view function decorated with @accepts/@returns, driven inside a request context.
|
|
"""
|
|
|
|
from collections.abc import Callable
|
|
from functools import wraps
|
|
from io import BytesIO
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
from flask import Flask, request
|
|
from flask.ctx import RequestContext
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
from sqlalchemy.orm import Session
|
|
from werkzeug.exceptions import RequestEntityTooLarge, UnprocessableEntity
|
|
|
|
from controllers.common.schema import register_response_schema_model, register_schema_model
|
|
from controllers.openapi import openapi_ns
|
|
from controllers.openapi._contract import (
|
|
MULTIPART_HEADER_BYTES_PER_PART,
|
|
CapRequestBody,
|
|
accepts,
|
|
multipart_header_allowance,
|
|
paginated,
|
|
returns,
|
|
)
|
|
from controllers.openapi._errors import RequestOverRouteLimits
|
|
from controllers.openapi._models import Hint, PaginationEnvelope
|
|
from controllers.openapi._upload import UploadPart
|
|
from tests.helpers.unit.controllers.openapi.auth._world import account_subject, make_ctx
|
|
|
|
|
|
class ContractQuery(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
page: int = Field(1, ge=1)
|
|
limit: int = Field(20, ge=1, le=100)
|
|
|
|
|
|
class ContractBody(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
name: str
|
|
|
|
|
|
class ContractUploadBody(BaseModel):
|
|
package: UploadPart
|
|
|
|
|
|
class ContractResp(BaseModel):
|
|
value: int
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="module")
|
|
def _register_contract_test_models():
|
|
# Register for @accepts(body=)/@returns name lookups; drop on teardown so these
|
|
# test-only models don't leak into the shared openapi_ns / generated spec.
|
|
register_schema_model(openapi_ns, ContractBody)
|
|
register_schema_model(openapi_ns, ContractUploadBody)
|
|
register_response_schema_model(openapi_ns, ContractResp)
|
|
yield
|
|
openapi_ns.models.pop(ContractBody.__name__, None)
|
|
openapi_ns.models.pop(ContractUploadBody.__name__, None)
|
|
openapi_ns.models.pop(ContractResp.__name__, None)
|
|
|
|
|
|
def _guard_like(view):
|
|
"""Stand-in for ``@subject_router.guard`` — an outermost @wraps layer."""
|
|
|
|
@wraps(view)
|
|
def wrapper(*args, **kwargs):
|
|
return view(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
class _Page(PaginationEnvelope[int]):
|
|
pass
|
|
|
|
|
|
def test_paginated_fills_the_next_page_hint_from_the_call_and_leaves_the_rest_alone():
|
|
@paginated("thing.list")
|
|
def view(*, ctx, workspace_id, query):
|
|
return _Page.build(page=query.page, limit=query.limit, total=50, items=[])
|
|
|
|
result = view(ctx=object(), workspace_id="ws-1", query=ContractQuery(page=1, limit=20))
|
|
assert result.hints == [
|
|
Hint(summary="Next page", op="thing.list", input={"workspace_id": "ws-1", "page": 2, "limit": 20})
|
|
]
|
|
assert view(ctx=object(), workspace_id="ws-1", query=ContractQuery(page=3, limit=20)).hints == []
|
|
|
|
@paginated("thing.get")
|
|
def plain(*, ctx):
|
|
return ContractResp(value=1)
|
|
|
|
assert plain(ctx=object()) == ContractResp(value=1)
|
|
|
|
|
|
def test_paginated_keeps_the_view_hints_and_adds_the_next_page() -> None:
|
|
setup = Hint(summary="Set up", op="thing.describe", input={})
|
|
|
|
@paginated("thing.list")
|
|
def view(*, ctx: object, query: ContractQuery) -> _Page:
|
|
page = _Page.build(page=query.page, limit=query.limit, total=50, items=[])
|
|
page.hints = [setup]
|
|
return page
|
|
|
|
assert [h.op for h in view(ctx=object(), query=ContractQuery(page=1, limit=20)).hints] == [
|
|
"thing.describe",
|
|
"thing.list",
|
|
]
|
|
|
|
own = Hint(summary="Next page", op="thing.list", input={"cursor": "c2"})
|
|
|
|
@paginated("thing.list")
|
|
def cursor_view(*, ctx: object, query: ContractQuery) -> _Page:
|
|
page = _Page.build(page=query.page, limit=query.limit, total=50, items=[])
|
|
page.hints = [own]
|
|
return page
|
|
|
|
assert cursor_view(ctx=object(), query=ContractQuery(page=1, limit=20)).hints == [own]
|
|
|
|
|
|
def test_accepts_injects_validated_query_with_defaults_for_absent_fields(app):
|
|
@accepts(query=ContractQuery)
|
|
def view(*, query):
|
|
return query
|
|
|
|
with app.test_request_context("/?page=3"):
|
|
result = view()
|
|
|
|
assert isinstance(result, ContractQuery)
|
|
assert result.page == 3
|
|
assert result.limit == 20
|
|
|
|
|
|
class ContractListQuery(BaseModel):
|
|
tag: list[str] | None = None
|
|
word: str | None = None
|
|
|
|
|
|
def test_accepts_fills_list_typed_query_fields_from_repeated_parameters(app: Flask) -> None:
|
|
@accepts(query=ContractListQuery)
|
|
def view(*, query: ContractListQuery) -> ContractListQuery:
|
|
return query
|
|
|
|
with app.test_request_context("/?tag=a&tag=b&word=x&word=y"):
|
|
result = view()
|
|
|
|
assert result.tag == ["a", "b"]
|
|
assert result.word == "x"
|
|
|
|
|
|
@pytest.mark.parametrize("query_string", ["page=abc", "unknown=1"])
|
|
def test_accepts_rejects_invalid_query_with_422(app, query_string):
|
|
@accepts(query=ContractQuery)
|
|
def view(*, query):
|
|
return query
|
|
|
|
with app.test_request_context(f"/?{query_string}"):
|
|
with pytest.raises(UnprocessableEntity):
|
|
view()
|
|
|
|
|
|
def test_accepts_validation_error_is_sanitized_and_structured(app):
|
|
"""422 body is structured and leaks neither the pydantic docs url nor the user input."""
|
|
|
|
@accepts(body=ContractBody)
|
|
def view(*, body):
|
|
return body
|
|
|
|
with app.test_request_context("/", method="POST", json={"secret": "leak-me"}):
|
|
with pytest.raises(UnprocessableEntity) as exc_info:
|
|
view()
|
|
|
|
data = cast(dict[str, Any], cast(Any, exc_info.value).data)
|
|
assert data["message"] == "Request validation failed"
|
|
assert isinstance(data["errors"], list)
|
|
assert data["errors"]
|
|
for err in data["errors"]:
|
|
assert {"type", "loc", "msg"} <= err.keys()
|
|
assert "url" not in err
|
|
assert "input" not in err
|
|
assert "leak-me" not in str(data)
|
|
|
|
|
|
def test_accepts_injects_validated_body(app):
|
|
@accepts(body=ContractBody)
|
|
def view(*, body):
|
|
return body
|
|
|
|
with app.test_request_context("/", method="POST", json={"name": "x"}):
|
|
result = view()
|
|
|
|
assert isinstance(result, ContractBody)
|
|
assert result.name == "x"
|
|
|
|
|
|
def test_returns_serializes_model_with_decorator_status(app):
|
|
@returns(200, ContractResp)
|
|
def view():
|
|
return ContractResp(value=7)
|
|
|
|
with app.test_request_context("/"):
|
|
body, status = view()
|
|
|
|
assert status == 200
|
|
assert body == {"value": 7}
|
|
|
|
|
|
@pytest.mark.parametrize("trailing", [(201,), (202, {"X-Test": "1"})], ids=["status", "status_and_headers"])
|
|
def test_returns_serializes_model_in_tuple_and_keeps_trailing_parts(app, trailing):
|
|
@returns(200, ContractResp)
|
|
def view():
|
|
return ContractResp(value=9), *trailing
|
|
|
|
with app.test_request_context("/"):
|
|
body, *rest = view()
|
|
|
|
assert body == {"value": 9}
|
|
assert tuple(rest) == trailing
|
|
|
|
|
|
def test_returns_passes_through_non_model(app):
|
|
sentinel = object()
|
|
|
|
@returns(200, ContractResp)
|
|
def view():
|
|
return sentinel
|
|
|
|
with app.test_request_context("/"):
|
|
result = view()
|
|
|
|
assert result is sentinel
|
|
|
|
|
|
# Swagger metadata (read off __apidoc__) must survive @wraps up through the guard layer.
|
|
|
|
|
|
def test_accepts_returns_emit_apidoc_through_guard_stack():
|
|
@_guard_like
|
|
@returns(200, ContractResp)
|
|
@accepts(query=ContractQuery, body=ContractBody)
|
|
def view(*, query, body):
|
|
return ContractResp(value=1)
|
|
|
|
apidoc = getattr(view, "__apidoc__", {})
|
|
assert "page" in apidoc.get("params", {}) # from @accepts(query=)
|
|
assert apidoc.get("expect") # from @accepts(body=), via @openapi_ns.expect
|
|
assert "200" in apidoc.get("responses", {}) # from @returns (flask_restx keys by str code)
|
|
|
|
|
|
_MIB = 1024 * 1024
|
|
|
|
|
|
def _post(app: Flask, size: int, **extra: str) -> RequestContext:
|
|
data = {"package": (BytesIO(b"x" * size), "a.yaml"), **extra}
|
|
return app.test_request_context("/x", method="POST", data=data, content_type="multipart/form-data")
|
|
|
|
|
|
@pytest.fixture
|
|
def view() -> Callable[[], str]:
|
|
@accepts(body=ContractUploadBody)
|
|
def handler(*, body: ContractUploadBody) -> str:
|
|
return "ok"
|
|
|
|
return handler
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("size", "extra", "accepted"),
|
|
[
|
|
(1, {"note": '"y"'}, True),
|
|
(2 * _MIB, {}, False),
|
|
(1, {"note": '"y"', "other": '"z"'}, False),
|
|
],
|
|
)
|
|
def test_a_capped_route_sets_its_own_limits_and_names_them(
|
|
app: Flask,
|
|
sqlite_session: Session,
|
|
view: Callable[[], str],
|
|
size: int,
|
|
extra: dict[str, str],
|
|
accepted: bool,
|
|
) -> None:
|
|
app.config["MAX_FORM_PARTS"] = 1
|
|
cap = CapRequestBody(max_bytes=lambda: _MIB, max_parts=lambda: 2)
|
|
with _post(app, size, **extra):
|
|
cap.run(account_subject(), make_ctx(sqlite_session), sqlite_session)
|
|
assert "form" not in request.__dict__
|
|
if accepted:
|
|
assert view() == "ok"
|
|
return
|
|
with pytest.raises(RequestOverRouteLimits, match="at most 2 parts and 1 MiB$"):
|
|
view()
|
|
|
|
|
|
def test_an_uncapped_route_keeps_the_generic_413(app: Flask, view: Callable[[], str]) -> None:
|
|
app.config["MAX_FORM_PARTS"] = 1
|
|
with _post(app, 1, extra="y"), pytest.raises(RequestEntityTooLarge) as exc:
|
|
view()
|
|
assert not isinstance(exc.value, RequestOverRouteLimits)
|
|
|
|
|
|
def test_multipart_header_allowance_scales_with_the_part_count() -> None:
|
|
assert multipart_header_allowance(0) == 0
|
|
assert multipart_header_allowance(5002) == 5002 * MULTIPART_HEADER_BYTES_PER_PART
|