cubefs

1
// Copyright 2022 The CubeFS Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12
// implied. See the License for the specific language governing
13
// permissions and limitations under the License.
14

15
package rpc
16

17
import (
18
	"errors"
19
	"fmt"
20
	"net/http"
21
	"testing"
22

23
	"github.com/stretchr/testify/require"
24
)
25

26
type (
27
	testMiddleware struct {
28
		Name   string
29
		RunIn  func() error
30
		RunOut func() error
31
		ProgressHandler
32
	}
33
)
34

35
func (mw testMiddleware) Handler(w http.ResponseWriter, req *http.Request, f func(http.ResponseWriter, *http.Request)) {
36
	fmt.Println("In  Middleware", mw.Name)
37
	if err := mw.RunIn(); err != nil {
38
		return
39
	}
40
	f(w, req)
41
	if err := mw.RunOut(); err != nil {
42
		return
43
	}
44
	fmt.Println("Out Middleware", mw.Name)
45
}
46

47
func TestMiddlewareBase(t *testing.T) {
48
	{
49
		var routed string
50
		router := New()
51
		router.Handle(http.MethodGet, "/", func(c *Context) { routed += "app" })
52
		handler := MiddlewareHandlerWith(router)
53

54
		w := new(mockResponseWriter)
55
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
56
		handler.ServeHTTP(w, req)
57
		require.Equal(t, "app", routed)
58
	}
59
	{
60
		var routed string
61
		router := New()
62
		router.Handle(http.MethodGet, "/", func(c *Context) { routed += "app" })
63
		handler := MiddlewareHandlerFuncWith(router,
64
			testMiddleware{
65
				Name: "middleware1",
66
				RunIn: func() error {
67
					routed += "1-in"
68
					return nil
69
				},
70
				RunOut: func() error {
71
					routed += "1-out"
72
					return nil
73
				},
74
			},
75
		)
76

77
		w := new(mockResponseWriter)
78
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
79
		handler.ServeHTTP(w, req)
80
		require.Equal(t, "1-inapp1-out", routed)
81
	}
82
	{
83
		var routed string
84
		router := New()
85
		router.Handle(http.MethodGet, "/", func(c *Context) { routed += "app" })
86
		handler := MiddlewareHandlerWith(router,
87
			testMiddleware{
88
				Name: "middleware1",
89
				RunIn: func() error {
90
					routed += "1-in"
91
					return errors.New("")
92
				},
93
				RunOut: func() error {
94
					routed += "1-out"
95
					return nil
96
				},
97
			},
98
		)
99

100
		w := new(mockResponseWriter)
101
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
102
		handler.ServeHTTP(w, req)
103
		require.Equal(t, "1-in", routed)
104
	}
105
	{
106
		var routed string
107
		router := New()
108
		router.Handle(http.MethodGet, "/", func(c *Context) { routed += "app" })
109

110
		handler := MiddlewareHandlerWith(router,
111
			testMiddleware{
112
				Name: "middleware1",
113
				RunIn: func() error {
114
					routed += "1-in"
115
					return nil
116
				},
117
				RunOut: func() error {
118
					routed += "1-out"
119
					return nil
120
				},
121
			},
122
			testMiddleware{
123
				Name: "middleware2",
124
				RunIn: func() error {
125
					routed += "2-in"
126
					return nil
127
				},
128
				RunOut: func() error {
129
					routed += "2-out"
130
					return errors.New("")
131
				},
132
			},
133
			testMiddleware{
134
				Name: "middleware3",
135
				RunIn: func() error {
136
					routed += "3-in"
137
					return errors.New("")
138
				},
139
				RunOut: func() error {
140
					routed += "3-out"
141
					return errors.New("")
142
				},
143
			},
144
			testMiddleware{
145
				Name: "middleware4",
146
				RunIn: func() error {
147
					routed += "4-in"
148
					return errors.New("")
149
				},
150
				RunOut: func() error {
151
					routed += "4-out"
152
					return nil
153
				},
154
			},
155
		)
156

157
		w := new(mockResponseWriter)
158
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
159
		handler.ServeHTTP(w, req)
160
		require.Equal(t, "1-in2-in3-in2-out1-out", routed)
161
	}
162
}
163

164
func TestMiddlewareDefault(t *testing.T) {
165
	defer initDefaultRouter()
166
	{
167
		var routed string
168
		GET("/", func(c *Context) { routed += "app" })
169
		handler := MiddlewareHandler()
170

171
		w := new(mockResponseWriter)
172
		req, _ := http.NewRequest(http.MethodGet, "/", nil)
173
		handler.ServeHTTP(w, req)
174
		require.Equal(t, "app", routed)
175
	}
176
	{
177
		var routed string
178
		GET("/get", func(c *Context) { routed += "app" })
179
		handler := MiddlewareHandlerFunc(
180
			testMiddleware{
181
				Name: "middleware1",
182
				RunIn: func() error {
183
					routed += "1-in"
184
					return nil
185
				},
186
				RunOut: func() error {
187
					routed += "1-out"
188
					return nil
189
				},
190
			},
191
		)
192

193
		w := new(mockResponseWriter)
194
		req, _ := http.NewRequest(http.MethodGet, "/get", nil)
195
		handler.ServeHTTP(w, req)
196
		require.Equal(t, "1-inapp1-out", routed)
197
	}
198
}
199