Skip to content

Commit 9e41bd1

Browse files
Raajheer1gyuho
Raajheer1
and
gyuho
authored
Add testing for pkg/server (#680)
Continued progress towards #472 --------- Co-authored-by: Gyuho Lee <[email protected]>
1 parent e550c29 commit 9e41bd1

File tree

3 files changed

+457
-0
lines changed

3 files changed

+457
-0
lines changed

‎pkg/server/handlers_test.go‎

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
package server
2+
3+
import (
4+
"encoding/json"
5+
"net/http"
6+
"net/http/httptest"
7+
"testing"
8+
"time"
9+
10+
"github.com/gin-gonic/gin"
11+
"github.com/stretchr/testify/assert"
12+
"sigs.k8s.io/yaml"
13+
14+
"github.com/leptonai/gpud/components"
15+
gpudconfig "github.com/leptonai/gpud/pkg/config"
16+
"github.com/leptonai/gpud/pkg/metrics"
17+
)
18+
19+
func TestNewGlobalHandler(t *testing.T) {
20+
var metricStore metrics.Store
21+
ghler := newGlobalHandler(nil, components.NewRegistry(nil), metricStore)
22+
assert.NotNil(t, ghler)
23+
}
24+
25+
func TestGetReqTime(t *testing.T) {
26+
g := &globalHandler{}
27+
gin.SetMode(gin.TestMode)
28+
29+
tests := []struct {
30+
name string
31+
startTimeQuery string
32+
endTimeQuery string
33+
expectError bool
34+
}{
35+
{
36+
name: "no query params",
37+
startTimeQuery: "",
38+
endTimeQuery: "",
39+
expectError: false,
40+
},
41+
{
42+
name: "valid start and end times",
43+
startTimeQuery: "1609459200", // 2021-01-01 00:00:00
44+
endTimeQuery: "1609545600", // 2021-01-02 00:00:00
45+
expectError: false,
46+
},
47+
{
48+
name: "invalid start time",
49+
startTimeQuery: "invalid",
50+
endTimeQuery: "1609545600",
51+
expectError: true,
52+
},
53+
{
54+
name: "invalid end time",
55+
startTimeQuery: "1609459200",
56+
endTimeQuery: "invalid",
57+
expectError: true,
58+
},
59+
}
60+
61+
for _, tt := range tests {
62+
t.Run(tt.name, func(t *testing.T) {
63+
w := httptest.NewRecorder()
64+
c, _ := gin.CreateTestContext(w)
65+
req := httptest.NewRequest("GET", "/?startTime="+tt.startTimeQuery+"&endTime="+tt.endTimeQuery, nil)
66+
c.Request = req
67+
68+
startTime, endTime, err := g.getReqTime(c)
69+
70+
if tt.expectError {
71+
assert.Error(t, err)
72+
} else {
73+
assert.NoError(t, err)
74+
if tt.startTimeQuery != "" {
75+
expectedStartTime := time.Unix(1609459200, 0)
76+
assert.Equal(t, expectedStartTime, startTime)
77+
}
78+
if tt.endTimeQuery != "" {
79+
expectedEndTime := time.Unix(1609545600, 0)
80+
assert.Equal(t, expectedEndTime, endTime)
81+
}
82+
}
83+
})
84+
}
85+
}
86+
87+
func TestCreateHealthzHandler(t *testing.T) {
88+
gin.SetMode(gin.TestMode)
89+
router := gin.New()
90+
router.GET("/healthz", createHealthzHandler())
91+
92+
tests := []struct {
93+
name string
94+
contentType string
95+
jsonIndent string
96+
wantStatus int
97+
checkBody func(t *testing.T, body []byte)
98+
}{
99+
{
100+
name: "default JSON response",
101+
wantStatus: http.StatusOK,
102+
checkBody: func(t *testing.T, body []byte) {
103+
var resp Healthz
104+
err := json.Unmarshal(body, &resp)
105+
assert.NoError(t, err)
106+
assert.Equal(t, "ok", resp.Status)
107+
assert.Equal(t, "v1", resp.Version)
108+
},
109+
},
110+
{
111+
name: "indented JSON response",
112+
jsonIndent: "true",
113+
wantStatus: http.StatusOK,
114+
checkBody: func(t *testing.T, body []byte) {
115+
var resp Healthz
116+
err := json.Unmarshal(body, &resp)
117+
assert.NoError(t, err)
118+
assert.Equal(t, "ok", resp.Status)
119+
assert.Equal(t, "v1", resp.Version)
120+
},
121+
},
122+
{
123+
name: "YAML response",
124+
contentType: "application/yaml",
125+
wantStatus: http.StatusOK,
126+
checkBody: func(t *testing.T, body []byte) {
127+
var resp Healthz
128+
err := yaml.Unmarshal(body, &resp)
129+
assert.NoError(t, err)
130+
assert.Equal(t, "ok", resp.Status)
131+
assert.Equal(t, "v1", resp.Version)
132+
},
133+
},
134+
}
135+
136+
for _, tt := range tests {
137+
t.Run(tt.name, func(t *testing.T) {
138+
req := httptest.NewRequest("GET", "/healthz", nil)
139+
if tt.contentType != "" {
140+
req.Header.Set("Content-Type", tt.contentType)
141+
}
142+
if tt.jsonIndent != "" {
143+
req.Header.Set("json-indent", tt.jsonIndent)
144+
}
145+
w := httptest.NewRecorder()
146+
router.ServeHTTP(w, req)
147+
148+
assert.Equal(t, tt.wantStatus, w.Code)
149+
if tt.checkBody != nil {
150+
tt.checkBody(t, w.Body.Bytes())
151+
}
152+
})
153+
}
154+
}
155+
156+
func TestCreateConfigHandler(t *testing.T) {
157+
gin.SetMode(gin.TestMode)
158+
159+
cfg := &gpudconfig.Config{
160+
Address: "localhost:8080",
161+
}
162+
163+
router := gin.New()
164+
router.GET("/config", createConfigHandler(cfg))
165+
166+
tests := []struct {
167+
name string
168+
contentType string
169+
jsonIndent string
170+
wantStatus int
171+
checkBody func(t *testing.T, body []byte)
172+
}{
173+
{
174+
name: "default JSON response",
175+
wantStatus: http.StatusOK,
176+
checkBody: func(t *testing.T, body []byte) {
177+
var resp gpudconfig.Config
178+
err := json.Unmarshal(body, &resp)
179+
assert.NoError(t, err)
180+
assert.Equal(t, "localhost:8080", resp.Address)
181+
},
182+
},
183+
{
184+
name: "indented JSON response",
185+
jsonIndent: "true",
186+
wantStatus: http.StatusOK,
187+
checkBody: func(t *testing.T, body []byte) {
188+
var resp gpudconfig.Config
189+
err := json.Unmarshal(body, &resp)
190+
assert.NoError(t, err)
191+
assert.Equal(t, "localhost:8080", resp.Address)
192+
},
193+
},
194+
{
195+
name: "YAML response",
196+
contentType: "application/yaml",
197+
wantStatus: http.StatusOK,
198+
checkBody: func(t *testing.T, body []byte) {
199+
var resp gpudconfig.Config
200+
err := yaml.Unmarshal(body, &resp)
201+
assert.NoError(t, err)
202+
assert.Equal(t, "localhost:8080", resp.Address)
203+
},
204+
},
205+
}
206+
207+
for _, tt := range tests {
208+
t.Run(tt.name, func(t *testing.T) {
209+
req := httptest.NewRequest("GET", "/config", nil)
210+
if tt.contentType != "" {
211+
req.Header.Set("Content-Type", tt.contentType)
212+
}
213+
if tt.jsonIndent != "" {
214+
req.Header.Set("json-indent", tt.jsonIndent)
215+
}
216+
w := httptest.NewRecorder()
217+
router.ServeHTTP(w, req)
218+
219+
assert.Equal(t, tt.wantStatus, w.Code)
220+
if tt.checkBody != nil {
221+
tt.checkBody(t, w.Body.Bytes())
222+
}
223+
})
224+
}
225+
}

‎pkg/server/middleware_test.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package server
2+
3+
import (
4+
"net/http"
5+
"net/http/httptest"
6+
"testing"
7+
8+
"github.com/gin-gonic/gin"
9+
"github.com/stretchr/testify/assert"
10+
"go.uber.org/zap/zaptest"
11+
)
12+
13+
func TestInstallRootGinMiddlewares(t *testing.T) {
14+
router := gin.New()
15+
installRootGinMiddlewares(router)
16+
17+
router.GET("/test", func(c *gin.Context) {
18+
c.String(http.StatusOK, "test")
19+
})
20+
21+
req := httptest.NewRequest("GET", "/test", nil)
22+
w := httptest.NewRecorder()
23+
24+
router.ServeHTTP(w, req)
25+
26+
// Check response
27+
assert.Equal(t, http.StatusOK, w.Code)
28+
assert.NotEmpty(t, w.Header().Get("X-Request-Id"), "X-Request-Id header should be set")
29+
30+
// Verify that ContextWithFallback was set
31+
assert.True(t, router.ContextWithFallback, "ContextWithFallback should be true")
32+
}
33+
34+
func TestInstallCommonGinMiddlewares(t *testing.T) {
35+
// Create a test logger
36+
logger := zaptest.NewLogger(t)
37+
router := gin.New()
38+
installCommonGinMiddlewares(router, logger)
39+
40+
router.GET("/test", func(c *gin.Context) {
41+
c.String(http.StatusOK, "test")
42+
})
43+
44+
req := httptest.NewRequest("GET", "/test", nil)
45+
w := httptest.NewRecorder()
46+
47+
router.ServeHTTP(w, req)
48+
49+
assert.Equal(t, http.StatusOK, w.Code)
50+
51+
router.GET("/panic", func(c *gin.Context) {
52+
panic("test panic")
53+
})
54+
55+
req = httptest.NewRequest("GET", "/panic", nil)
56+
w = httptest.NewRecorder()
57+
58+
router.ServeHTTP(w, req)
59+
60+
// Check that we got a 500 error but the server didn't crash
61+
assert.Equal(t, http.StatusInternalServerError, w.Code)
62+
}

0 commit comments

Comments
 (0)