pkg/expect: ExpectFunc, LineCount

ExpectFunc to make expect more extensible. LineCount to be
able to check 'no output' command.
This commit is contained in:
Gyu-Ho Lee 2016-04-06 15:43:44 -07:00
parent 62990fb5fa
commit 679e5e379b
2 changed files with 57 additions and 3 deletions

View File

@ -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.
func (ep *ExpectProcess) Expect(s string) (string, error) {
// ExpectFunc returns the first line satisfying the function f.
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) }

View File

@ -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 {