From b91ad5d48571a9c86f328229715552e6e8f58f71 Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Fri, 3 Apr 2026 16:06:00 +0000
Subject: [PATCH] feat(ansible): support wait_for sleep

Co-Authored-By: Virgil <virgil@lethean.io>
---
 modules.go          | 19 ++++++++++++-------
 modules_adv_test.go | 46 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/modules.go b/modules.go
index 1a144d1..15a9686 100644
--- a/modules.go
+++ b/modules.go
@@ -2888,7 +2888,12 @@ func (e *Executor) moduleWaitFor(ctx context.Context, client sshExecutorClient,
 	searchRegex := getStringArg(args, "search_regex", "")
 	timeoutMsg := getStringArg(args, "msg", "wait_for timed out")
 	delay := getIntArg(args, "delay", 0)
+	sleep := getIntArg(args, "sleep", 1)
 	timeout := getIntArg(args, "timeout", 300)
+	pollInterval := time.Duration(sleep) * time.Second
+	if pollInterval <= 0 {
+		pollInterval = 250 * time.Millisecond
+	}
 	var compiledRegex *regexp.Regexp
 	if searchRegex != "" {
 		var err error
@@ -2910,7 +2915,7 @@ func (e *Executor) moduleWaitFor(ctx context.Context, client sshExecutorClient,
 
 	if path != "" {
 		deadline := time.NewTimer(time.Duration(timeout) * time.Second)
-		ticker := time.NewTicker(250 * time.Millisecond)
+		ticker := time.NewTicker(pollInterval)
 		defer deadline.Stop()
 		defer ticker.Stop()
 
@@ -2958,24 +2963,24 @@ func (e *Executor) moduleWaitFor(ctx context.Context, client sshExecutorClient,
 	if port > 0 {
 		switch state {
 		case "started", "present":
-			cmd := sprintf("timeout %d bash -c 'until nc -z %s %d; do sleep 1; done'",
-				timeout, host, port)
+			cmd := sprintf("timeout %d bash -c 'until nc -z %s %d; do sleep %d; done'",
+				timeout, host, port, sleep)
 			stdout, stderr, rc, err := client.Run(ctx, cmd)
 			if err != nil || rc != 0 {
 				return &TaskResult{Failed: true, Msg: stderr, Stdout: stdout, RC: rc}, nil
 			}
 			return &TaskResult{Changed: false}, nil
 		case "stopped", "absent":
-			cmd := sprintf("timeout %d bash -c 'until ! nc -z %s %d; do sleep 1; done'",
-				timeout, host, port)
+			cmd := sprintf("timeout %d bash -c 'until ! nc -z %s %d; do sleep %d; done'",
+				timeout, host, port, sleep)
 			stdout, stderr, rc, err := client.Run(ctx, cmd)
 			if err != nil || rc != 0 {
 				return &TaskResult{Failed: true, Msg: stderr, Stdout: stdout, RC: rc}, nil
 			}
 			return &TaskResult{Changed: false}, nil
 		case "drained":
-			cmd := sprintf("timeout %d bash -c 'until ! ss -Htan state established \"( sport = :%d or dport = :%d )\" | grep -q .; do sleep 1; done'",
-				timeout, port, port)
+			cmd := sprintf("timeout %d bash -c 'until ! ss -Htan state established \"( sport = :%d or dport = :%d )\" | grep -q .; do sleep %d; done'",
+				timeout, port, port, sleep)
 			stdout, stderr, rc, err := client.Run(ctx, cmd)
 			if err != nil || rc != 0 {
 				return &TaskResult{Failed: true, Msg: stderr, Stdout: stdout, RC: rc}, nil
diff --git a/modules_adv_test.go b/modules_adv_test.go
index 5cb73f8..dbc8b6d 100644
--- a/modules_adv_test.go
+++ b/modules_adv_test.go
@@ -1246,6 +1246,52 @@ func TestModulesAdv_ModuleWaitFor_Good_WaitsForPortDrained(t *testing.T) {
 	assert.True(t, mock.hasExecuted(`ss -Htan state established`))
 }
 
+func TestModulesAdv_ModuleWaitFor_Good_UsesCustomSleepInterval(t *testing.T) {
+	e, mock := newTestExecutorWithMock("host1")
+	mock.addFile("/tmp/slow-ready", []byte("ready=false\n"))
+
+	go func() {
+		time.Sleep(150 * time.Millisecond)
+		mock.mu.Lock()
+		mock.files["/tmp/slow-ready"] = []byte("ready=true\n")
+		mock.mu.Unlock()
+	}()
+
+	start := time.Now()
+	result, err := e.moduleWaitFor(context.Background(), mock, map[string]any{
+		"path":         "/tmp/slow-ready",
+		"search_regex": "ready=true",
+		"sleep":        2,
+		"timeout":      3,
+	})
+	elapsed := time.Since(start)
+
+	require.NoError(t, err)
+	assert.NotNil(t, result)
+	assert.False(t, result.Failed)
+	assert.False(t, result.Changed)
+	assert.GreaterOrEqual(t, elapsed, 2*time.Second)
+}
+
+func TestModulesAdv_ModuleWaitFor_Good_UsesCustomSleepInPortLoop(t *testing.T) {
+	e, mock := newTestExecutorWithMock("host1")
+	mock.expectCommand(`timeout 5 bash -c 'until nc -z 127.0.0.1 8080; do sleep 3; done'`, "", "", 0)
+
+	result, err := e.moduleWaitFor(context.Background(), mock, map[string]any{
+		"host":    "127.0.0.1",
+		"port":    8080,
+		"state":   "started",
+		"timeout": 5,
+		"sleep":   3,
+	})
+
+	require.NoError(t, err)
+	assert.NotNil(t, result)
+	assert.False(t, result.Failed)
+	assert.False(t, result.Changed)
+	assert.True(t, mock.hasExecuted(`sleep 3`))
+}
+
 func TestModulesAdv_ModuleWaitFor_Good_AcceptsStringNumericArgs(t *testing.T) {
 	e, mock := newTestExecutorWithMock("host1")
 	mock.expectCommand(`timeout 0 bash -c 'until ! nc -z 127.0.0.1 8080; do sleep 1; done'`, "", "", 0)