testutils: add ExpectExactNTimes function to LogObserver

Signed-off-by: Clement <[email protected]>

tests/framework/testutils/log_observer.go

@@ -46,8 +46,8 @@ func NewLogObserver(level zapcore.LevelEnabler) (zapcore.Core, *LogObserver) {
 	}
 }
 
-// Expect returns the first N lines containing the given string.
-func (logOb *LogObserver) Expect(ctx context.Context, s string, count int) ([]string, error) {
+// ExpectAtLeastNTimes returns the first N lines containing the given string.
+func (logOb *LogObserver) ExpectAtLeastNTimes(ctx context.Context, s string, count int) ([]string, error) {
 	return logOb.ExpectFunc(ctx, func(log string) bool { return strings.Contains(log, s) }, count)
 }
 
@@ -87,6 +87,60 @@ func (logOb *LogObserver) ExpectFunc(ctx context.Context, filter func(string) bo
 	}
 }
 
+// ExpectExactNTimes returns all lines that satisfy the filter if there are
+// exactly `count` of them when the duration ends. Otherwise, it returns an error.
+// Make sure ctx has a timeout longer than duration for ExpectExactNTimes to work properly.
+func (logOb *LogObserver) ExpectExactNTimes(ctx context.Context, s string, count int, duration time.Duration) ([]string, error) {
+	return logOb.ExpectExactNTimesFunc(ctx, func(log string) bool { return strings.Contains(log, s) }, count, duration)
+}
+
+// ExpectExactNTimesFunc returns all lines that satisfy the filter if there are
+// exactly `count` of them when the duration ends. Otherwise, it returns an error.
+// Make sure ctx has a timeout longer than duration for ExpectExactNTimesFunc to work properly.
+func (logOb *LogObserver) ExpectExactNTimesFunc(ctx context.Context, filter func(string) bool, count int, duration time.Duration) ([]string, error) {
+	timer := time.NewTimer(duration)
+	defer timer.Stop()
+
+	i := 0
+	res := make([]string, 0, count)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-timer.C:
+			if len(res) == count {
+				return res, nil
+			} else {
+				return nil, fmt.Errorf("failed to expect, expected: %d, got: %d", count, len(res))
+			}
+		default:
+		}
+
+		entries := logOb.syncLogs()
+
+		// The order of entries won't be changed because of append-only.
+		// It's safe to skip scanned entries by reusing `i`.
+		for ; i < len(entries); i++ {
+			buf, err := logOb.enc.EncodeEntry(entries[i].Entry, entries[i].Context)
+			if err != nil {
+				return nil, fmt.Errorf("failed to encode entry: %w", err)
+			}
+
+			logInStr := buf.String()
+			if filter(logInStr) {
+				res = append(res, logInStr)
+			}
+
+			if len(res) > count {
+				return nil, fmt.Errorf("failed to expect; too many occurrences; expected: %d, got:%d", count, len(res))
+			}
+		}
+
+		time.Sleep(10 * time.Millisecond)
+	}
+}
+
 // syncLogs is to take all the existing logged entries from zapobserver and
 // truncate zapobserver's entries slice.
 func (logOb *LogObserver) syncLogs() []zapobserver.LoggedEntry {

tests/framework/testutils/log_observer_test.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -34,7 +35,7 @@ func TestLogObserver_Timeout(t *testing.T) {
 	logger.Info(t.Name())
 
 	ctx, cancel := context.WithTimeout(context.TODO(), 100*time.Millisecond)
-	_, err := logOb.Expect(ctx, "unknown", 1)
+	_, err := logOb.ExpectAtLeastNTimes(ctx, "unknown", 1)
 	cancel()
 	assert.True(t, errors.Is(err, context.DeadlineExceeded))
 
@@ -53,7 +54,7 @@ func TestLogObserver_Expect(t *testing.T) {
 	go func() {
 		defer close(resCh)
 
-		res, err := logOb.Expect(ctx, t.Name(), 2)
+		res, err := logOb.ExpectAtLeastNTimes(ctx, t.Name(), 2)
 		require.NoError(t, err)
 		resCh <- res
 	}()
@@ -81,3 +82,76 @@ func TestLogObserver_Expect(t *testing.T) {
 
 	assert.Equal(t, 2, len(logOb.entries))
 }
+
+// TestLogObserver_ExpectExactNTimes tests the behavior of ExpectExactNTimes.
+// It covers 4 scenarios:
+//
+// 1. Occurrences in log is less than expected when duration ends.
+// 2. Occurrences in log is exactly as expected when duration ends.
+// 3. Occurrences in log is greater than expected when duration ends.
+// 4. context deadline exceeded before duration ends
+func TestLogObserver_ExpectExactNTimes(t *testing.T) {
+	logCore, logOb := NewLogObserver(zap.InfoLevel)
+
+	logger := zap.New(logCore)
+
+	// Log messages at 0ms, 100ms, 200ms
+	msgs := []string{"Hello " + t.Name(), t.Name() + ", World", t.Name()}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
+	defer cancel()
+
+	wg := sync.WaitGroup{}
+
+	// When duration ends at 50ms, there's only one message in the log.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		_, err := logOb.ExpectExactNTimes(ctx, t.Name(), 2, 50*time.Millisecond)
+		assert.ErrorContains(t, err, "failed to expect, expected: 2, got: 1")
+	}()
+
+	// When duration ends at 150ms, there are 2 messages in the log.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		lines, err := logOb.ExpectExactNTimes(ctx, t.Name(), 2, 150*time.Millisecond)
+		assert.NoError(t, err)
+		assert.Len(t, lines, 2)
+		for i := 0; i < 2; i++ {
+			// lines should be like:
+			// 2023-04-16T11:46:19.367+0800 INFO	Hello TestLogObserver_ExpectExactNTimes
+			// 2023-04-16T11:46:19.408+0800	INFO	TestLogObserver_ExpectExactNTimes, World
+			//
+			// msgs:
+			// Hello TestLogObserver_ExpectExactNTimes
+			// TestLogObserver_ExpectExactNTimes, World
+			// TestLogObserver_ExpectExactNTimes
+			strings.HasSuffix(lines[i], msgs[i])
+		}
+	}()
+
+	// Before duration ends at 250ms, there are already 3 messages in the log.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		_, err := logOb.ExpectExactNTimes(ctx, t.Name(), 2, 250*time.Millisecond)
+		assert.ErrorContains(t, err,"failed to expect; too many occurrences; expected: 2, got:3")
+	}()
+
+	// context deadline exceeded at 300ms before the 400ms duration can end.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		_, err := logOb.ExpectExactNTimes(ctx, t.Name(), 3, 400*time.Millisecond)
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+	}()
+
+	// Log messages at 0ms, 100ms, 200ms
+	for _, msg := range msgs {
+		logger.Info(msg)
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	wg.Wait()
+}

tests/integration/v3_watch_restore_test.go

@@ -177,7 +177,7 @@ func expectMemberLog(t *testing.T, m *integration.Member, timeout time.Duration,
 	ctx, cancel := context.WithTimeout(context.TODO(), timeout)
 	defer cancel()
 
-	lines, err := m.LogObserver.Expect(ctx, s, count)
+	lines, err := m.LogObserver.ExpectAtLeastNTimes(ctx, s, count)
 	if err != nil {
 		t.Fatalf("failed to expect (log:%s, count:%v): %v", s, count, err)
 	}