test: add test

zeromicro · Aug 27, 2024 · a01f817 · a01f817
1 parent 7d33874
commit a01f817
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/rest/engine_test.go b/rest/engine_test.go
@@ -114,7 +114,8 @@ Verbose: true
 		{
 			priority: true,
 			jwt: jwtSetting{
-				enabled: true,
+				enabled:   true,
+				tokenKeys: []string{"Token", "X-Token"},
 			},
 			signature: signatureSetting{},
 			routes: []Route{{

diff --git a/rest/handler/authhandler_test.go b/rest/handler/authhandler_test.go
@@ -57,6 +57,32 @@ func TestAuthHandler(t *testing.T) {
 	assert.Equal(t, "content", resp.Body.String())
 }
 
+func TestAuthHandler_WithTokenKeys(t *testing.T) {
+	const key = "B63F477D-BBA3-4E52-96D3-C0034C27694A"
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
+	token, err := buildToken(key, map[string]any{
+		"key": "value",
+	}, 3600)
+	assert.Nil(t, err)
+	req.Header.Set("X-Token", token)
+	handler := Authorize(key, WithTokenKeys([]string{"Token", "X-Token"}))(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("X-Test", "test")
+			_, err := w.Write([]byte("content"))
+			assert.Nil(t, err)
+
+			flusher, ok := w.(http.Flusher)
+			assert.True(t, ok)
+			flusher.Flush()
+		}))
+
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, http.StatusOK, resp.Code)
+	assert.Equal(t, "test", resp.Header().Get("X-Test"))
+	assert.Equal(t, "content", resp.Body.String())
+}
+
 func TestAuthHandlerWithPrevSecret(t *testing.T) {
 	const (
 		key     = "14F17379-EB8F-411B-8F12-6929002DCA76"