veyron/lib/expect: update ExpectSet(Eventually)?RE

This updates ExpectSet(Eventually)?RE to return the matches as soon as
they are seen, rather than after EOF or a timeout occurs.

All unmatched regular expressions are tested against each line as the
line is read. As soon as there are no more unmatched regular
expressions, the matches are returned and no more output is consumed.

Change-Id: If4fc823a9fb47ae47a1968835726d2e68cd1045e
diff --git a/lib/expect/expect.go b/lib/expect/expect.go
index 01e9e35..8388d99 100644
--- a/lib/expect/expect.go
+++ b/lib/expect/expect.go
@@ -237,7 +237,7 @@
 }
 
 // ExpectRE asserts that the next line in the input matches the pattern using
-// regexp.MustCompile(pattern,n).FindAllStringSubmatch.
+// regexp.MustCompile(pattern).FindAllStringSubmatch(..., n).
 func (s *Session) ExpectRE(pattern string, n int) [][]string {
 	if s.Failed() {
 		return [][]string{}
@@ -279,7 +279,9 @@
 // patterns in the order that they are supplied as parameters. Consequently
 // the set may contain repetitions if the same pattern is expected multiple
 // times. The value returned is either:
-//   * nil in the case of an error or no match, or
+//   * nil in the case of an error, or
+//   * nil if n lines are read or EOF is encountered before all expressions are
+//       matched, or
 //   * an array of length len(expected), whose ith element contains the result
 //       of FindStringSubmatch of expected[i] on the matching string (never
 //       nil). If there are no capturing groups in expected[i], the return
@@ -296,15 +298,17 @@
 	}
 }
 
-// ExpectSetEventuallyRE is like ExpectSetRE except that it reads
-// all remaining output rather than just the next n lines and thus
-// can be used to look for a set of patterns that occur within that
-// output. The value returned is either:
-//   * nil in the case of an error or no match, or
+// ExpectSetEventuallyRE is like ExpectSetRE except that it reads as much
+// output as required rather than just the next n lines. The value returned is
+// either:
+//   * nil in the case of an error, or
+//   * nil if EOF is encountered before all expressions are matched, or
 //   * an array of length len(expected), whose ith element contains the result
 //       of FindStringSubmatch of expected[i] on the matching string (never
 //       nil). If there are no capturing groups in expected[i], the return
 //       value's [i][0] will contain the entire matching string
+// This function stops consuming output as soon as all regular expressions are
+// matched.
 func (s *Session) ExpectSetEventuallyRE(expected ...string) [][]string {
 	if s.Failed() {
 		return nil
@@ -318,9 +322,10 @@
 }
 
 // expectSetRE will look for the expected set of patterns in the next
-// numLines of output or in all remaining output.
+// numLines of output or in all remaining output. If all expressions are
+// matched, no more output is consumed.
 func (s *Session) expectSetRE(numLines int, expected ...string) ([][]string, error) {
-
+	matches := make([][]string, len(expected))
 	regexps := make([]*regexp.Regexp, len(expected))
 	for i, expRE := range expected {
 		re, err := regexp.Compile(expRE)
@@ -329,9 +334,12 @@
 		}
 		regexps[i] = re
 	}
-	actual := []string{}
 	i := 0
+	matchCount := 0
 	for {
+		if matchCount == len(expected) {
+			break
+		}
 		line, err := s.read(readLine)
 		line = strings.TrimRight(line, "\n")
 		s.log(err, "ExpectSetRE: %s", line)
@@ -341,34 +349,38 @@
 			}
 			break
 		}
-		actual = append(actual, line)
+
+		// Match the line against all regexp's and remove each regexp
+		// that matches.
+		for i, re := range regexps {
+			if re == nil {
+				continue
+			}
+			match := re.FindStringSubmatch(line)
+			if match != nil {
+				matchCount++
+				regexps[i] = nil
+				matches[i] = match
+				// Don't allow this line to be matched by more than one re.
+				break
+			}
+		}
+
 		i++
 		if numLines > 0 && i >= numLines {
 			break
 		}
 	}
 
-	matches := make([][]string, len(expected))
-	// Match each line against all regexp's and remove each regexp
-	// that matches.
-	for _, l := range actual {
-		for i, re := range regexps {
-			if re == nil {
-				continue
-			}
-			match := re.FindStringSubmatch(l)
-			if match != nil {
-				regexps[i] = nil
-				matches[i] = match
-				break
-			}
+	// It's an error if there are any unmatched regexps.
+	unmatchedRes := make([]string, 0)
+	for i, re := range regexps {
+		if re != nil {
+			unmatchedRes = append(unmatchedRes, expected[i])
 		}
 	}
-	// It's an error if there are any unmatched regexps.
-	for _, re := range regexps {
-		if re != nil {
-			return nil, fmt.Errorf("found no match for %q", re)
-		}
+	if len(unmatchedRes) > 0 {
+		return nil, fmt.Errorf("found no match for [%v]", strings.Join(unmatchedRes, ","))
 	}
 	return matches, nil
 }
diff --git a/lib/expect/expect_test.go b/lib/expect/expect_test.go
index b0fc48b..fb6f579 100644
--- a/lib/expect/expect_test.go
+++ b/lib/expect/expect_test.go
@@ -31,7 +31,6 @@
 	} else {
 		t.Log(s.Error())
 	}
-	s.ExpectEOF()
 }
 
 func TestExpectf(t *testing.T) {
@@ -40,10 +39,10 @@
 	buffer.WriteString("bar 22\n")
 	s := expect.NewSession(nil, bufio.NewReader(buffer), time.Minute)
 	s.Expectf("bar %d", 22)
+	s.ExpectEOF()
 	if err := s.Error(); err != nil {
 		t.Error(err)
 	}
-	s.ExpectEOF()
 }
 
 func TestEOF(t *testing.T) {
@@ -84,7 +83,6 @@
 			t.Errorf("missing or wrong error: %v", s.Error())
 		}
 	}
-	s.ExpectEOF()
 }
 
 func TestExpectSetRE(t *testing.T) {
@@ -101,15 +99,14 @@
 	}
 	want := [][]string{{"bar=baz"}, {"def"}, {"abc"}, {"abc"}}
 	if !reflect.DeepEqual(got, want) {
-		t.Error("unexpected result from ExpectSetRE, got %v, want %v", got, want)
+		t.Errorf("unexpected result from ExpectSetRE, got %v, want %v", got, want)
 	}
 	buffer.WriteString("ooh\n")
 	buffer.WriteString("aah\n")
 	s.ExpectSetRE("bar=.*", "def")
-	if got, want := s.Error(), "expect_test.go:108: found no match for \"bar=.*\""; got == nil || got.Error() != want {
-		t.Errorf("got %v, want %q", got, want)
+	if got, want := s.Error(), "found no match for [bar=.*,def]"; got == nil || !strings.Contains(got.Error(), want) {
+		t.Errorf("got %v, wanted something containing %q", got, want)
 	}
-	s.ExpectEOF()
 
 	buf = []byte{}
 	buffer = bytes.NewBuffer(buf)
@@ -120,6 +117,63 @@
 	if want := [][]string{{"hello world", "world"}, {"this is a test", "is", "a"}}; !reflect.DeepEqual(want, matches) {
 		t.Errorf("unexpected result from ExpectSetRE, got %v, want %v", matches, want)
 	}
+
+	buf = []byte{}
+	buffer = bytes.NewBuffer(buf)
+	s = expect.NewSession(nil, bufio.NewReader(buffer), time.Minute)
+	buffer.WriteString("aaa\n")
+	buffer.WriteString("aaa\n")
+	buffer.WriteString("aaa\n")
+
+	// Expect 3 x aaa to match.
+	s.ExpectSetRE("aaa", "aaa", "aaa")
+	if s.Error() != nil {
+		t.Errorf("unexpected error: %v", s.Error())
+	}
+
+	// Expecting one more aaa should fail: the entire input should have been consumed.
+	s.ExpectSetRE("aaa")
+	if s.Error() == nil {
+		t.Errorf("expected error but got none")
+	}
+
+	// Test a buffer that contains a match but not within the number of lines we expect.
+	buf = []byte{}
+	buffer = bytes.NewBuffer(buf)
+	buffer.WriteString("aaa\n")
+	buffer.WriteString("bbb\n")
+	s = expect.NewSession(nil, bufio.NewReader(buffer), time.Minute)
+	s.ExpectSetRE("bbb")
+	if s.Error() == nil {
+		t.Fatalf("expected error but got none")
+	}
+
+	// Test a buffer that contains a match and leaves us with nothing more to read.
+	buf = []byte{}
+	buffer = bytes.NewBuffer(buf)
+	buffer.WriteString("aaa\n")
+	buffer.WriteString("bbb\n")
+	s = expect.NewSession(nil, bufio.NewReader(buffer), time.Minute)
+	s.ExpectSetRE("bbb")
+	if s.Error() == nil {
+		t.Fatalf("expected error but got none")
+	}
+
+	// Now ensure that each regular expression matches a unique line.
+	buf = []byte{}
+	buffer = bytes.NewBuffer(buf)
+	buffer.WriteString("a 1\n")
+	buffer.WriteString("a 2\n")
+	buffer.WriteString("a 3\n")
+	s = expect.NewSession(nil, bufio.NewReader(buffer), time.Minute)
+	matches = s.ExpectSetRE("\\w (\\d)", "a (\\d)", "a (\\d)")
+	want = [][]string{{"a 1", "1"}, {"a 2", "2"}, {"a 3", "3"}}
+	if !reflect.DeepEqual(matches, want) {
+		t.Fatalf("unexpected result from ExpectSetRE, got %v, want %v", matches, want)
+	}
+	if s.ExpectEOF() != nil {
+		t.Fatalf("expected EOF but did not get it")
+	}
 }
 
 func TestExpectSetEventuallyRE(t *testing.T) {
@@ -135,10 +189,19 @@
 	if s.Error() != nil {
 		t.Errorf("unexpected error: %s", s.Error())
 	}
+
+	// Should see one more abc match after the we read def.
 	s.ExpectSetEventuallyRE("abc")
-	if got, want := s.Error(), "expect_test.go:138: found no match for \"abc\""; got == nil || got.Error() != want {
-		t.Errorf("got %q, want %q", got, want)
+	if s.Error() != nil {
+		t.Errorf("unexpected error: %s", s.Error())
 	}
+
+	// Trying to match abc again should yield an error.
+	s.ExpectSetEventuallyRE("abc")
+	if got, want := s.Error(), "found no match for [abc]"; got == nil || !strings.Contains(got.Error(), want) {
+		t.Errorf("got %q, wanted something containing %q", got, want)
+	}
+
 	// Need to clear the EOF from the previous ExpectSetEventuallyRE call
 	buf = []byte{}
 	buffer = bytes.NewBuffer(buf)
@@ -146,10 +209,9 @@
 	buffer.WriteString("ooh\n")
 	buffer.WriteString("aah\n")
 	s.ExpectSetEventuallyRE("zzz")
-	if got, want := s.Error(), "expect_test.go:148: found no match for \"zzz\""; got == nil || got.Error() != want {
-		t.Errorf("got %q, want %q", got, want)
+	if got, want := s.Error(), "found no match for [zzz]"; got == nil || !strings.Contains(got.Error(), want) {
+		t.Errorf("got %q, wanted something containing %q", got, want)
 	}
-	s.ExpectEOF()
 
 	buf = []byte{}
 	buffer = bytes.NewBuffer(buf)
@@ -161,6 +223,18 @@
 	if want := [][]string{{"hello world", "world"}, {"this is a test", "is", "a"}}; !reflect.DeepEqual(want, matches) {
 		t.Errorf("unexpected result from ExpectSetRE, got %v, want %v", matches, want)
 	}
+
+	// Test error output with multiple unmatched res.
+	buf = []byte{}
+	buffer = bytes.NewBuffer(buf)
+	s = expect.NewSession(nil, bufio.NewReader(buffer), time.Minute)
+	buffer.WriteString("not expected\n")
+	buffer.WriteString("hello world\n")
+	buffer.WriteString("this is a test\n")
+	s.ExpectSetEventuallyRE("blargh", "blerg", "blorg")
+	if got, want := s.Error(), "found no match for [blargh,blerg,blorg]"; !strings.Contains(got.Error(), want) {
+		t.Errorf("got %q, wanted something containing %q", got, want)
+	}
 }
 
 func TestRead(t *testing.T) {
@@ -192,5 +266,7 @@
 	if got != want {
 		t.Errorf("got %q, want %q", got, want)
 	}
-	s.ExpectEOF()
+	if s.ExpectEOF() != nil {
+		t.Fatalf("expected EOF but did not get it")
+	}
 }