github.com/kubewharf/katalyst-core@v0.5.3/pkg/util/process/http_test.go (about) 1 /* 2 Copyright 2022 The Katalyst Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package process 18 19 import ( 20 "context" 21 "fmt" 22 "net/http" 23 "net/url" 24 "testing" 25 "time" 26 27 "github.com/stretchr/testify/assert" 28 29 "github.com/kubewharf/katalyst-core/pkg/metrics" 30 ) 31 32 type dummyHandler struct { 33 success int 34 } 35 36 func (d *dummyHandler) ServeHTTP(_ http.ResponseWriter, _ *http.Request) { 37 d.success++ 38 } 39 40 type dummyResponseWriter struct{} 41 42 func (d dummyResponseWriter) Header() http.Header { return make(http.Header) } 43 func (d dummyResponseWriter) Write([]byte) (int, error) { return 0, nil } 44 func (d dummyResponseWriter) WriteHeader(_ int) {} 45 46 func TestHTTPHandler(t *testing.T) { 47 t.Parallel() 48 49 httpCleanupVisitorPeriod = time.Second 50 51 type req struct { 52 burst int 53 success int 54 remoteAddr string 55 } 56 57 for _, tc := range []struct { 58 comment string 59 enabled []string 60 reqs []req 61 visitCnt int 62 }{ 63 { 64 comment: "empty chain to pass all", 65 enabled: []string{}, 66 reqs: []req{ 67 { 68 burst: 2, 69 success: 2, 70 remoteAddr: "addr-1", 71 }, 72 }, 73 visitCnt: 0, 74 }, 75 { 76 comment: "rate limiter chain to limit requests", 77 enabled: []string{HTTPChainRateLimiter}, 78 reqs: []req{ 79 { 80 burst: 1, 81 success: 1, 82 remoteAddr: "addr-1", 83 }, 84 { 85 burst: 2, 86 success: 1, 87 remoteAddr: "addr-2", 88 }, 89 }, 90 visitCnt: 2, 91 }, 92 { 93 comment: "auth chain to build auth for user", 94 enabled: []string{HTTPChainCredential}, 95 reqs: []req{ 96 { 97 burst: 2, 98 success: 2, 99 remoteAddr: "addr-1", 100 }, 101 { 102 burst: 1, 103 success: 1, 104 remoteAddr: "addr-2", 105 }, 106 }, 107 visitCnt: 0, 108 }, 109 { 110 comment: "auth chain to limit requests", 111 enabled: []string{HTTPChainCredential, HTTPChainRateLimiter}, 112 reqs: []req{ 113 { 114 burst: 2, 115 success: 1, 116 remoteAddr: "addr-1", 117 }, 118 { 119 burst: 2, 120 success: 0, 121 remoteAddr: "addr-2", 122 }, 123 }, 124 visitCnt: 1, 125 }, 126 } { 127 t.Logf("test case: %v", tc.comment) 128 h := NewHTTPHandler(tc.enabled, []string{}, true, metrics.DummyMetrics{}) 129 130 ctx, cancel := context.WithCancel(context.Background()) 131 h.Run(ctx) 132 time.Sleep(time.Second) 133 134 for _, r := range tc.reqs { 135 f := &dummyHandler{} 136 hf := h.WithHandleChain(f) 137 138 for j := 0; j < r.burst; j++ { 139 hr := &http.Request{ 140 Header: make(http.Header), 141 URL: &url.URL{ 142 Path: "fake_path", 143 }, 144 RemoteAddr: fmt.Sprintf("%v", r.remoteAddr), 145 Method: fmt.Sprintf("%v-order-%v", r.remoteAddr, j), 146 } 147 148 hf.ServeHTTP(dummyResponseWriter{}, hr) 149 } 150 151 assert.Equal(t, r.success, f.success) 152 } 153 154 h.mux.Lock() 155 assert.Equal(t, tc.visitCnt, len(h.visitors)) 156 h.mux.Unlock() 157 158 time.Sleep(time.Second * 3) 159 160 h.mux.Lock() 161 assert.Equal(t, 0, len(h.visitors)) 162 h.mux.Unlock() 163 164 cancel() 165 } 166 }