1 package pprof
2
3 import (
4 "net/http"
5 "net/http/httptest"
6 "testing"
7
8 "github.com/gin-gonic/gin"
9 )
10
11 func Test_getPrefix(t *testing.T) {
12 tests := []struct {
13 name string
14 args []string
15 want string
16 }{
17 {"default value", nil, "/debug/pprof"},
18 {"test user input value", []string{"test/pprof"}, "test/pprof"},
19 {"test user input value", []string{"test/pprof", "pprof"}, "test/pprof"},
20 }
21 for _, tt := range tests {
22 if got := getPrefix(tt.args...); got != tt.want {
23 t.Errorf("%q. getPrefix() = %v, want %v", tt.name, got, tt.want)
24 }
25 }
26 }
27
28 func TestRegisterAndRouteRegister(t *testing.T) {
29 bearerToken := "Bearer token"
30 gin.SetMode(gin.ReleaseMode)
31 r := gin.New()
32 Register(r)
33 adminGroup := r.Group("/admin", func(c *gin.Context) {
34 if c.Request.Header.Get("Authorization") != bearerToken {
35 c.AbortWithStatus(http.StatusForbidden)
36 return
37 }
38 c.Next()
39 })
40 RouteRegister(adminGroup, "pprof")
41
42 req, _ := http.NewRequest(http.MethodGet, "/debug/pprof/", nil)
43 rw := httptest.NewRecorder()
44 r.ServeHTTP(rw, req)
45
46 if expected, got := http.StatusOK, rw.Code; expected != got {
47 t.Errorf("expected: %d, got: %d", expected, got)
48 }
49
50 req, _ = http.NewRequest(http.MethodGet, "/admin/pprof/", nil)
51 rw = httptest.NewRecorder()
52 r.ServeHTTP(rw, req)
53
54 if expected, got := http.StatusForbidden, rw.Code; expected != got {
55 t.Errorf("expected: %d, got: %d", expected, got)
56 }
57
58 req, _ = http.NewRequest(http.MethodGet, "/admin/pprof/", nil)
59 req.Header.Set("Authorization", bearerToken)
60 rw = httptest.NewRecorder()
61 r.ServeHTTP(rw, req)
62
63 if expected, got := http.StatusOK, rw.Code; expected != got {
64 t.Errorf("expected: %d, got: %d", expected, got)
65 }
66 }
67
View as plain text