fix: allways init response header map at first. (#1066)

swaggo · Nov 29, 2021 · 231fa42 · 231fa42
1 parent 305a5cf
commit 231fa42
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 3 deletions.
diff --git a/operation.go b/operation.go
@@ -887,6 +887,14 @@ func (operation *Operation) ParseResponseHeaderComment(commentLine string, _ *as
 		}
 	}
 
+	for code, response := range operation.Responses.StatusCodeResponses {
+		if response.Headers == nil {
+			r := operation.Responses.StatusCodeResponses[code]
+			r.Headers = make(map[string]spec.Header)
+			operation.Responses.StatusCodeResponses[code] = r
+		}
+	}
+
 	if strings.EqualFold(matches[1], "all") {
 		if operation.Responses.Default != nil {
 			operation.Responses.Default.Headers[headerKey] = header
@@ -918,9 +926,6 @@ func (operation *Operation) ParseResponseHeaderComment(commentLine string, _ *as
 		if operation.Responses.StatusCodeResponses != nil {
 			response, responseExist := operation.Responses.StatusCodeResponses[code]
 			if responseExist {
-				if response.Headers == nil {
-					response.Headers = make(map[string]spec.Header)
-				}
 				response.Headers[headerKey] = header
 
 				operation.Responses.StatusCodeResponses[code] = response

diff --git a/operation_test.go b/operation_test.go
@@ -997,6 +997,60 @@ func TestParseResponseCommentWithHeaderForCodes(t *testing.T) {
 	assert.Error(t, err, "ParseComment should not fail")
 }
 
+func TestParseResponseCommentWithHeaderOnlyAll(t *testing.T) {
+	t.Parallel()
+
+	operation := NewOperation(nil)
+
+	comment := `@Success 200,201,default "it's ok"`
+	err := operation.ParseComment(comment, nil)
+	assert.NoError(t, err, "ParseComment should not fail")
+
+	comment = `@Header all {string} Token "qwerty"`
+	err = operation.ParseComment(comment, nil)
+	assert.NoError(t, err, "ParseComment should not fail")
+
+	b, err := json.MarshalIndent(operation, "", "    ")
+	assert.NoError(t, err)
+
+	expected := `{
+    "responses": {
+        "200": {
+            "description": "it's ok",
+            "headers": {
+                "Token": {
+                    "type": "string",
+                    "description": "qwerty"
+                }
+            }
+        },
+        "201": {
+            "description": "it's ok",
+            "headers": {
+                "Token": {
+                    "type": "string",
+                    "description": "qwerty"
+                }
+            }
+        },
+        "default": {
+            "description": "it's ok",
+            "headers": {
+                "Token": {
+                    "type": "string",
+                    "description": "qwerty"
+                }
+            }
+        }
+    }
+}`
+	assert.Equal(t, expected, string(b))
+
+	comment = `@Header 200 "Mallformed"`
+	err = operation.ParseComment(comment, nil)
+	assert.Error(t, err, "ParseComment should not fail")
+}
+
 func TestParseEmptyResponseOnlyCode(t *testing.T) {
 	t.Parallel()