Add content unit option to customize referencable entities (#130)

swaggest · Sep 9, 2024 · cdbf0a7 · cdbf0a7
1 parent 6ab9f42
commit cdbf0a7
Show file tree

Hide file tree

Showing 7 changed files with 193 additions and 5 deletions.
diff --git a/openapi3/helper.go b/openapi3/helper.go
@@ -210,3 +210,15 @@ func (s *Spec) SetHTTPBearerTokenSecurity(securityName string, format string, de
 		},
 	)
 }
+
+// SetReference sets a reference and discards existing content.
+func (r *ResponseOrRef) SetReference(ref string) {
+	r.ResponseReferenceEns().Ref = ref
+	r.Response = nil
+}
+
+// SetReference sets a reference and discards existing content.
+func (r *RequestBodyOrRef) SetReference(ref string) {
+	r.RequestBodyReferenceEns().Ref = ref
+	r.RequestBody = nil
+}
diff --git a/openapi3/reflect.go b/openapi3/reflect.go
@@ -300,6 +300,10 @@ func (r *Reflector) setupRequest(o *Operation, oc openapi.OperationContext) erro
 		if cu.Description != "" && o.RequestBody != nil && o.RequestBody.RequestBody != nil {
 			o.RequestBody.RequestBody.WithDescription(cu.Description)
 		}
+
+		if cu.Customize != nil && o.RequestBody != nil {
+			cu.Customize(o.RequestBody)
+		}
 	}
 
 	return nil
@@ -668,10 +672,16 @@ func (r *Reflector) setupResponse(o *Operation, oc openapi.OperationContext) err
 			resp.Description = http.StatusText(cu.HTTPStatus)
 		}
 
+		ror := ResponseOrRef{Response: resp}
+
+		if cu.Customize != nil {
+			cu.Customize(&ror)
+		}
+
 		if cu.IsDefault {
-			o.Responses.Default = &ResponseOrRef{Response: resp}
+			o.Responses.Default = &ror
 		} else {
-			o.Responses.WithMapOfResponseOrRefValuesItem(httpStatus, ResponseOrRef{Response: resp})
+			o.Responses.WithMapOfResponseOrRefValuesItem(httpStatus, ror)
 		}
 	}
 

diff --git a/openapi3/reflect_test.go b/openapi3/reflect_test.go
@@ -1298,3 +1298,60 @@ func TestNewReflector_examples(t *testing.T) {
 	  }
 	}`, r.SpecSchema())
 }
+
+func TestWithCustomize(t *testing.T) {
+	r := openapi3.NewReflector()
+
+	op, err := r.NewOperationContext(http.MethodPost, "/{document_id}/{client}")
+	require.NoError(t, err)
+
+	op.AddReqStructure(new(struct {
+		DocumentID string `path:"document_id"`
+		Client     string `path:"client"`
+		Foo        int    `json:"foo"`
+	}), openapi.WithCustomize(func(cor openapi.ContentOrReference) {
+		_, ok := cor.(*openapi3.RequestBodyOrRef)
+		assert.True(t, ok)
+
+		cor.SetReference("../somewhere/components/requests/foo.yaml")
+	}))
+
+	op.AddRespStructure(
+		nil, openapi.WithReference("../somewhere/components/responses/204.yaml"), openapi.WithHTTPStatus(204),
+	)
+	op.AddRespStructure(
+		nil, openapi.WithCustomize(func(cor openapi.ContentOrReference) {
+			_, ok := cor.(*openapi3.ResponseOrRef)
+			assert.True(t, ok)
+
+			cor.SetReference("../somewhere/components/responses/200.yaml")
+		}), openapi.WithHTTPStatus(200),
+	)
+
+	require.NoError(t, r.AddOperation(op))
+
+	assertjson.EqMarshal(t, `{
+	  "openapi":"3.0.3","info":{"title":"","version":""},
+	  "paths":{
+		"/{document_id}/{client}":{
+		  "post":{
+			"parameters":[
+			  {
+				"name":"document_id","in":"path","required":true,
+				"schema":{"type":"string"}
+			  },
+			  {
+				"name":"client","in":"path","required":true,
+				"schema":{"type":"string"}
+			  }
+			],
+			"requestBody":{"$ref":"../somewhere/components/requests/foo.yaml"},
+			"responses":{
+			  "200":{"$ref":"../somewhere/components/responses/200.yaml"},
+			  "204":{"$ref":"../somewhere/components/responses/204.yaml"}
+			}
+		  }
+		}
+	  }
+	}`, r.SpecSchema())
+}
diff --git a/openapi31/helper.go b/openapi31/helper.go
@@ -269,3 +269,15 @@ func (s *Spec) SetHTTPBearerTokenSecurity(securityName string, format string, de
 		},
 	)
 }
+
+// SetReference sets a reference and discards existing content.
+func (r *ResponseOrReference) SetReference(ref string) {
+	r.ReferenceEns().Ref = ref
+	r.Response = nil
+}
+
+// SetReference sets a reference and discards existing content.
+func (r *RequestBodyOrReference) SetReference(ref string) {
+	r.ReferenceEns().Ref = ref
+	r.RequestBody = nil
+}
diff --git a/openapi31/reflect.go b/openapi31/reflect.go
@@ -247,6 +247,10 @@ func (r *Reflector) setupRequest(o *Operation, oc openapi.OperationContext) erro
 		if cu.Description != "" && o.RequestBody != nil && o.RequestBody.RequestBody != nil {
 			o.RequestBody.RequestBody.WithDescription(cu.Description)
 		}
+
+		if cu.Customize != nil && o.RequestBody != nil {
+			cu.Customize(o.RequestBody)
+		}
 	}
 
 	return nil
@@ -619,10 +623,16 @@ func (r *Reflector) setupResponse(o *Operation, oc openapi.OperationContext) err
 			resp.Description = http.StatusText(cu.HTTPStatus)
 		}
 
+		ror := ResponseOrReference{Response: resp}
+
+		if cu.Customize != nil {
+			cu.Customize(&ror)
+		}
+
 		if cu.IsDefault {
-			o.Responses.Default = &ResponseOrReference{Response: resp}
+			o.Responses.Default = &ror
 		} else {
-			o.Responses.WithMapOfResponseOrReferenceValuesItem(httpStatus, ResponseOrReference{Response: resp})
+			o.Responses.WithMapOfResponseOrReferenceValuesItem(httpStatus, ror)
 		}
 	}
 

diff --git a/openapi31/reflect_test.go b/openapi31/reflect_test.go
@@ -1426,3 +1426,60 @@ func TestNewReflector_examples(t *testing.T) {
 	  }
 	}`, r.SpecSchema())
 }
+
+func TestWithCustomize(t *testing.T) {
+	r := openapi31.NewReflector()
+
+	op, err := r.NewOperationContext(http.MethodPost, "/{document_id}/{client}")
+	require.NoError(t, err)
+
+	op.AddReqStructure(new(struct {
+		DocumentID string `path:"document_id"`
+		Client     string `path:"client"`
+		Foo        int    `json:"foo"`
+	}), openapi.WithCustomize(func(cor openapi.ContentOrReference) {
+		_, ok := cor.(*openapi31.RequestBodyOrReference)
+		assert.True(t, ok)
+
+		cor.SetReference("../somewhere/components/requests/foo.yaml")
+	}))
+
+	op.AddRespStructure(
+		nil, openapi.WithReference("../somewhere/components/responses/204.yaml"), openapi.WithHTTPStatus(204),
+	)
+	op.AddRespStructure(
+		nil, openapi.WithCustomize(func(cor openapi.ContentOrReference) {
+			_, ok := cor.(*openapi31.ResponseOrReference)
+			assert.True(t, ok)
+
+			cor.SetReference("../somewhere/components/responses/200.yaml")
+		}), openapi.WithHTTPStatus(200),
+	)
+
+	require.NoError(t, r.AddOperation(op))
+
+	assertjson.EqMarshal(t, `{
+	  "openapi":"3.1.0","info":{"title":"","version":""},
+	  "paths":{
+		"/{document_id}/{client}":{
+		  "post":{
+			"parameters":[
+			  {
+				"name":"document_id","in":"path","required":true,
+				"schema":{"type":"string"}
+			  },
+			  {
+				"name":"client","in":"path","required":true,
+				"schema":{"type":"string"}
+			  }
+			],
+			"requestBody":{"$ref":"../somewhere/components/requests/foo.yaml"},
+			"responses":{
+			  "200":{"$ref":"../somewhere/components/responses/200.yaml"},
+			  "204":{"$ref":"../somewhere/components/responses/204.yaml"}
+			}
+		  }
+		}
+	  }
+	}`, r.SpecSchema())
+}
diff --git a/operation.go b/operation.go
@@ -37,10 +37,40 @@ type ContentUnit struct {
 	// IsDefault indicates default response.
 	IsDefault bool
 
-	Description  string
+	Description string
+
+	// Customize allows fine control over prepared content entities.
+	// The cor value can be asserted to one of these types:
+	// *openapi3.RequestBodyOrRef
+	// *openapi3.ResponseOrRef
+	// *openapi31.RequestBodyOrReference
+	// *openapi31.ResponseOrReference
+	Customize func(cor ContentOrReference)
+
 	fieldMapping map[In]map[string]string
 }
 
+// ContentOrReference defines content entity that can be a reference.
+type ContentOrReference interface {
+	SetReference(ref string)
+}
+
+// WithCustomize is a ContentUnit option.
+func WithCustomize(customize func(cor ContentOrReference)) ContentOption {
+	return func(cu *ContentUnit) {
+		cu.Customize = customize
+	}
+}
+
+// WithReference is a ContentUnit option.
+func WithReference(ref string) ContentOption {
+	return func(cu *ContentUnit) {
+		cu.Customize = func(cor ContentOrReference) {
+			cor.SetReference(ref)
+		}
+	}
+}
+
 // ContentUnitPreparer defines self-contained ContentUnit.
 type ContentUnitPreparer interface {
 	SetupContentUnit(cu *ContentUnit)