feat: set custom request body schemas

This commit is contained in:
Daniel G. Taylor 2021-03-11 10:52:22 -08:00
parent 933d5cb53e
commit 6cb2154568
No known key found for this signature in database
GPG key ID: 8D100732CA686E06
2 changed files with 36 additions and 3 deletions

View file

@ -135,6 +135,12 @@ func (o *Operation) NoBodyReadTimeout() {
o.bodyReadTimeout = 0
}
// RequestSchema allows overriding the generated input body schema, giving you
// more control over documentation and validation.
func (o *Operation) RequestSchema(s *schema.Schema) {
o.requestSchema = s
}
// Run registers the handler function for this operation. It should be of the
// form: `func (ctx huma.Context)` or `func (ctx huma.Context, input)` where
// input is your input struct describing the input parameters and/or body.
@ -182,9 +188,12 @@ func (o *Operation) Run(handler interface{}) {
// Get body if present.
if body, ok := input.FieldByName("Body"); ok {
o.requestModel = body.Type
o.requestSchema, err = schema.GenerateWithMode(body.Type, schema.ModeWrite, nil)
if err != nil {
panic(fmt.Errorf("unable to generate JSON schema: %w", err))
if o.requestSchema == nil {
o.requestSchema, err = schema.GenerateWithMode(body.Type, schema.ModeWrite, nil)
if err != nil {
panic(fmt.Errorf("unable to generate JSON schema: %w", err))
}
}
}

View file

@ -11,6 +11,7 @@ import (
"testing"
"time"
"github.com/istreamlabs/huma/schema"
"github.com/stretchr/testify/assert"
)
@ -338,3 +339,26 @@ func TestRouterAutoConfig(t *testing.T) {
},
})
}
func TestCustomRequestSchema(t *testing.T) {
app := newTestRouter()
op := app.Resource("/foo").Post("id", "doc", NewResponse(http.StatusOK, "ok"))
op.RequestSchema(&schema.Schema{
Type: schema.TypeInteger,
Minimum: schema.F(0),
Maximum: schema.F(100),
})
op.Run(func(ctx Context, input struct {
Body int
}) {
// Custom schema should be used so we never get here.
ctx.WriteHeader(http.StatusOK)
})
w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodPost, "/foo", strings.NewReader("1234"))
app.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Result().StatusCode)
}