From 8663704d6e7e90addde4156cccde113b8aff7a8f Mon Sep 17 00:00:00 2001
From: ycdesu <sychewo@gmail.com>
Date: Sat, 6 Feb 2021 10:43:25 +0800
Subject: [PATCH] util: create IsJSON/IsHTML type helper

---
 pkg/util/http_response.go      | 16 ++++++++++++++++
 pkg/util/http_response_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/pkg/util/http_response.go b/pkg/util/http_response.go
index 7eef0346b..76b683716 100644
--- a/pkg/util/http_response.go
+++ b/pkg/util/http_response.go
@@ -40,3 +40,19 @@ func (r *Response) DecodeJSON(o interface{}) error {
 func (r *Response) IsError() bool {
 	return r.StatusCode >= 400
 }
+
+func (r *Response) IsJSON() bool {
+	switch r.Header.Get("content-type") {
+	case "text/json", "application/json", "application/json; charset=utf-8":
+		return true
+	}
+	return false
+}
+
+func (r *Response) IsHTML() bool {
+	switch r.Header.Get("content-type") {
+	case "text/html":
+		return true
+	}
+	return false
+}
diff --git a/pkg/util/http_response_test.go b/pkg/util/http_response_test.go
index d2b1a4c13..af864f825 100644
--- a/pkg/util/http_response_test.go
+++ b/pkg/util/http_response_test.go
@@ -42,3 +42,33 @@ func TestResponse_IsError(t *testing.T) {
 		assert.Equal(t, isErr, resp.IsError())
 	}
 }
+
+func TestResponse_IsJSON(t *testing.T) {
+	cases := map[string]bool{
+		"text/json":                       true,
+		"application/json":                true,
+		"application/json; charset=utf-8": true,
+		"text/html":                       false,
+	}
+	for k, v := range cases {
+		resp := &Response{Response: &http.Response{}}
+		resp.Header = http.Header{}
+		resp.Header.Set("content-type", k)
+		assert.Equal(t, v, resp.IsJSON())
+	}
+}
+
+func TestResponse_IsHTML(t *testing.T) {
+	cases := map[string]bool{
+		"text/json":                       false,
+		"application/json":                false,
+		"application/json; charset=utf-8": false,
+		"text/html":                       true,
+	}
+	for k, v := range cases {
+		resp := &Response{Response: &http.Response{}}
+		resp.Header = http.Header{}
+		resp.Header.Set("content-type", k)
+		assert.Equal(t, v, resp.IsHTML())
+	}
+}