From 8921873e7492290ffa94505297d4b11b0ea53c16 Mon Sep 17 00:00:00 2001
From: Tobias <tobiaskochjp@gmail.com>
Date: Wed, 24 Jul 2024 10:03:54 +0900
Subject: [PATCH] add wait function

---
 request.go        |  2 ++
 testpilot.go      | 16 ++++++++++++++++
 testpilot_test.go |  2 ++
 3 files changed, 20 insertions(+)

diff --git a/request.go b/request.go
index 2014824..a39666c 100644
--- a/request.go
+++ b/request.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 )
 
 type Request struct {
@@ -20,6 +21,7 @@ type Request struct {
 	plan    *TestPlan
 	expect  *Expect
 	key     string
+	d       time.Duration
 }
 
 // Body sets the request body
diff --git a/testpilot.go b/testpilot.go
index 80c012d..b6fc3a8 100644
--- a/testpilot.go
+++ b/testpilot.go
@@ -11,11 +11,14 @@ import (
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 )
 
 var placeholderRegex = regexp.MustCompile(`{\.?([\w\.\d]+)}`)
 var findPlaceholderRegex = regexp.MustCompile(`{[^}]+}`)
 
+const waitMethod = "TESTPILOT_WAIT"
+
 type TestPlan struct {
 	t                *testing.T
 	name             string
@@ -50,6 +53,10 @@ func (p *TestPlan) Run() {
 	p.runCalled = true
 	p.t.Run(p.name, func(t *testing.T) {
 		for _, request := range p.requests {
+			if request.method == waitMethod {
+				time.Sleep(request.d)
+				continue
+			}
 			url, err := normalizeUrl(request.url, p.lastResponseBody, p.responseStore)
 			if err != nil {
 				t.Errorf(err.Error())
@@ -74,6 +81,15 @@ func (p *TestPlan) Request(method, url string) *Request {
 	return request
 }
 
+// Wait pauses the test plan for a given duration
+func (p *TestPlan) Wait(d time.Duration) {
+	request := &Request{
+		method: waitMethod,
+		d:      d,
+	}
+	p.requests = append(p.requests, request)
+}
+
 // Response returns the last response body
 func (p *TestPlan) Response() []byte {
 	return p.lastResponseBody
diff --git a/testpilot_test.go b/testpilot_test.go
index f44b590..b95387c 100644
--- a/testpilot_test.go
+++ b/testpilot_test.go
@@ -7,6 +7,7 @@ import (
 	"net/http/httptest"
 	"strconv"
 	"testing"
+	"time"
 )
 
 func Test_TestPlan(t *testing.T) {
@@ -39,6 +40,7 @@ func Test_TestPlan(t *testing.T) {
 	p.Request("GET", server.URL+"/{user.id}").Expect().Status(200).Body(ResponseComparer(p))
 	p.Request("GET", server.URL+"/{user.id}").Expect().Status(200).Body(AssertExists("id"))
 	p.Request("GET", server.URL+"/2").Expect().Status(404)
+	p.Wait(20 * time.Second)
 	p.Request("GET", server.URL+"/ping").Expect().Status(200).Body(AssertEqual("pong"))
 	p.Run()