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  }