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 + } + } +}