61 lines
1.6 KiB
Go
61 lines
1.6 KiB
Go
|
package middleware_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
"github.com/stretchr/testify/assert"
|
||
|
|
|
||
|
|
"veza-backend-api/internal/middleware"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestTimeoutMiddleware_PassesGivenEnoughTime(t *testing.T) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
router := gin.New()
|
||
|
|
router.Use(middleware.Timeout(100 * time.Millisecond))
|
||
|
|
|
||
|
|
router.GET("/slow", func(c *gin.Context) {
|
||
|
|
time.Sleep(10 * time.Millisecond)
|
||
|
|
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||
|
|
})
|
||
|
|
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
req := httptest.NewRequest("GET", "/slow", nil)
|
||
|
|
router.ServeHTTP(w, req)
|
||
|
|
|
||
|
|
assert.Equal(t, http.StatusOK, w.Code)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTimeoutMiddleware_ContextTimesOut(t *testing.T) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
router := gin.New()
|
||
|
|
// Set a very short timeout
|
||
|
|
router.Use(middleware.Timeout(10 * time.Millisecond))
|
||
|
|
|
||
|
|
router.GET("/too-slow", func(c *gin.Context) {
|
||
|
|
// Simulate work that takes longer than timeout
|
||
|
|
// Check context deadline
|
||
|
|
ctx := c.Request.Context()
|
||
|
|
select {
|
||
|
|
case <-time.After(50 * time.Millisecond):
|
||
|
|
c.JSON(http.StatusOK, gin.H{"status": "should not reach here"})
|
||
|
|
case <-ctx.Done():
|
||
|
|
// Context should be cancelled
|
||
|
|
// We don't write anything, middleware might handle it or the client hangs up
|
||
|
|
}
|
||
|
|
})
|
||
|
|
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
req := httptest.NewRequest("GET", "/too-slow", nil)
|
||
|
|
router.ServeHTTP(w, req)
|
||
|
|
|
||
|
|
// Since our simple implementation just sets context but tries an opportunistic 504:
|
||
|
|
// If the handler respects context, it returns early.
|
||
|
|
// If standard flow, we expect 504 if middleware catches it, or just connection close.
|
||
|
|
// Let's check if we get 504.
|
||
|
|
assert.Equal(t, http.StatusGatewayTimeout, w.Code)
|
||
|
|
}
|