From c153eb7aef91876c8ad1617401e942b8cc69c220 Mon Sep 17 00:00:00 2001
From: Arthur ANDRIEU <aandrieu@scaleway.com>
Date: Thu, 1 Feb 2024 13:44:14 +0100
Subject: [PATCH] feat(delay): add ability to delay with milliseconds

---
 pkg/api/delay.go  | 19 ++++++++++++++++---
 pkg/api/server.go |  2 +-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/pkg/api/delay.go b/pkg/api/delay.go
index 6655cf754..0af531172 100644
--- a/pkg/api/delay.go
+++ b/pkg/api/delay.go
@@ -3,6 +3,7 @@ package api
 import (
 	"math/rand"
 	"net/http"
+	"regexp"
 
 	"strconv"
 	"time"
@@ -58,13 +59,25 @@ func (s *Server) delayHandler(w http.ResponseWriter, r *http.Request) {
 
 	vars := mux.Vars(r)
 
-	delay, err := strconv.Atoi(vars["wait"])
+	wait := vars["wait"]
+	re := regexp.MustCompile(`(?P<duration>\d+)(?P<unit>ms|s)?`)
+
+	match := re.FindStringSubmatch(wait)
+	delay, err := strconv.Atoi(match[1])
+
+	unit := time.Second
+	unitString := "s"
+	if match[2] == "ms" {
+		unit = time.Millisecond
+		unitString = "ms"
+	}
+
 	if err != nil {
 		s.ErrorResponse(w, r, span, err.Error(), http.StatusBadRequest)
 		return
 	}
 
-	time.Sleep(time.Duration(delay) * time.Second)
+	time.Sleep(time.Duration(delay) * unit)
 
-	s.JSONResponse(w, r, map[string]int{"delay": delay})
+	s.JSONResponse(w, r, map[string]any{"delay": delay, "unit": unitString})
 }
diff --git a/pkg/api/server.go b/pkg/api/server.go
index b4224f45b..778b36de8 100644
--- a/pkg/api/server.go
+++ b/pkg/api/server.go
@@ -104,7 +104,7 @@ func (s *Server) registerHandlers() {
 	s.router.PathPrefix("/echo/").HandlerFunc(s.echoHandler)
 	s.router.HandleFunc("/env", s.envHandler).Methods("GET", "POST")
 	s.router.HandleFunc("/headers", s.echoHeadersHandler).Methods("GET", "POST")
-	s.router.HandleFunc("/delay/{wait:[0-9]+}", s.delayHandler).Methods("GET").Name("delay")
+	s.router.HandleFunc("/delay/{wait:[0-9]+(?:s|ms)?}", s.delayHandler).Methods("GET").Name("delay")
 	s.router.HandleFunc("/healthz", s.healthzHandler).Methods("GET")
 	s.router.HandleFunc("/readyz", s.readyzHandler).Methods("GET")
 	s.router.HandleFunc("/readyz/enable", s.enableReadyHandler).Methods("POST")