From 9bca54b14e27d18ac613c1cfd43a02e8588f60a1 Mon Sep 17 00:00:00 2001
From: HAYAMA_Kaoru <iyahaya@nifty.com>
Date: Fri, 29 Nov 2024 12:42:43 +0900
Subject: [PATCH] Add `(*Watcher) ExpectWithTimeout`

---
 main.go    |  5 ++++-
 watcher.go | 41 +++++++++++++++++++++++++++++++++--------
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/main.go b/main.go
index 66a958e..f42c75c 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"time"
 
 	"golang.org/x/term"
 
@@ -20,8 +21,10 @@ func loop(ptmx pty.Pty) error {
 		return err
 	}
 
-	_ = watcher.Expect("100")
+	i := watcher.ExpectWithTimeout(time.Duration(10*time.Second), "100")
+	// i := watcher.Expect("100")
 	io.WriteString(ptmx, "exit\r")
+	println(i)
 	return sh.Wait()
 }
 
diff --git a/watcher.go b/watcher.go
index 6fa7103..793a229 100644
--- a/watcher.go
+++ b/watcher.go
@@ -4,6 +4,7 @@ import (
 	"io"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/aymanbagabas/go-pty"
 )
@@ -44,16 +45,40 @@ func (W *Watcher) updateLastLine() {
 	}
 }
 
+func (W *Watcher) checkWords(token string, words []string) int {
+	W.lastline += token
+	for i, word := range words {
+		if strings.Contains(W.lastline, word) {
+			W.updateLastLine()
+			return i
+		}
+	}
+	W.updateLastLine()
+	return -1
+
+}
+
 func (W *Watcher) Expect(words ...string) int {
-	for frag := range W.ch {
-		W.lastline += frag
-		for i, word := range words {
-			if strings.Contains(W.lastline, word) {
-				W.updateLastLine()
-				return i
-			}
+	for token := range W.ch {
+		if found := W.checkWords(token, words); found >= 0 {
+			return found
 		}
-		W.updateLastLine()
 	}
 	return -1
 }
+
+func (W *Watcher) ExpectWithTimeout(d time.Duration, words ...string) int {
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+
+	for {
+		select {
+		case token := <-W.ch:
+			if found := W.checkWords(token, words); found >= 0 {
+				return found
+			}
+		case <-timer.C:
+			return -1
+		}
+	}
+}