pkg/expect: ExpectFunc, LineCount

ExpectFunc to make expect more extensible. LineCount to be
able to check 'no output' command.

pkg/expect/expect.go

@@ -35,6 +35,7 @@ type ExpectProcess struct {
 	cond  *sync.Cond // for broadcasting updates are avaiable
 	mu    sync.Mutex // protects lines and err
 	lines []string
+	count int // increment whenever new line gets added
 	err   error
 }
 
@@ -65,6 +66,7 @@ func (ep *ExpectProcess) read() {
 		ep.err = rerr
 		if l != "" {
 			ep.lines = append(ep.lines, l)
+			ep.count++
 			if len(ep.lines) == 1 {
 				ep.cond.Signal()
 			}
@@ -74,8 +76,8 @@ func (ep *ExpectProcess) read() {
 	ep.cond.Signal()
 }
 
-// Expect returns the first line containing the given string.
+// ExpectFunc returns the first line satisfying the function f.
-func (ep *ExpectProcess) Expect(s string) (string, error) {
+func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) {
 	ep.mu.Lock()
 	for {
 		for len(ep.lines) == 0 && ep.err == nil {
@@ -86,7 +88,7 @@ func (ep *ExpectProcess) Expect(s string) (string, error) {
 		}
 		l := ep.lines[0]
 		ep.lines = ep.lines[1:]
-		if strings.Contains(l, s) {
+		if f(l) {
 			ep.mu.Unlock()
 			return l, nil
 		}
@@ -95,6 +97,19 @@ func (ep *ExpectProcess) Expect(s string) (string, error) {
 	return "", ep.err
 }
 
+// Expect returns the first line containing the given string.
+func (ep *ExpectProcess) Expect(s string) (string, error) {
+	return ep.ExpectFunc(func(txt string) bool { return strings.Contains(txt, s) })
+}
+
+// LineCount returns the number of recorded lines since
+// the beginning of the process.
+func (ep *ExpectProcess) LineCount() int {
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
+	return ep.count
+}
+
 // Stop kills the expect process and waits for it to exit.
 func (ep *ExpectProcess) Stop() error { return ep.close(true) }
 

pkg/expect/expect_test.go

@@ -18,6 +18,24 @@ package expect
 
 import "testing"
 
+func TestExpectFunc(t *testing.T) {
+	ep, err := NewExpect("/bin/echo", "hello world")
+	if err != nil {
+		t.Fatal(err)
+	}
+	wstr := "hello world\r\n"
+	l, eerr := ep.ExpectFunc(func(a string) bool { return len(a) > 10 })
+	if eerr != nil {
+		t.Fatal(eerr)
+	}
+	if l != wstr {
+		t.Fatalf(`got "%v", expected "%v"`, l, wstr)
+	}
+	if cerr := ep.Close(); cerr != nil {
+		t.Fatal(cerr)
+	}
+}
+
 func TestEcho(t *testing.T) {
 	ep, err := NewExpect("/bin/echo", "hello world")
 	if err != nil {
@@ -39,6 +57,27 @@ func TestEcho(t *testing.T) {
 	}
 }
 
+func TestLineCount(t *testing.T) {
+	ep, err := NewExpect("/usr/bin/printf", "1\n2\n3")
+	if err != nil {
+		t.Fatal(err)
+	}
+	wstr := "3"
+	l, eerr := ep.Expect(wstr)
+	if eerr != nil {
+		t.Fatal(eerr)
+	}
+	if l != wstr {
+		t.Fatalf(`got "%v", expected "%v"`, l, wstr)
+	}
+	if ep.LineCount() != 3 {
+		t.Fatalf("got %d, expected 3", ep.LineCount())
+	}
+	if cerr := ep.Close(); cerr != nil {
+		t.Fatal(cerr)
+	}
+}
+
 func TestSend(t *testing.T) {
 	ep, err := NewExpect("/usr/bin/tr", "a", "b")
 	if err != nil {