acquisition tests: do not rely on hardcoded timeout to finish tests (#4610)

Author: blotus

pkg/acquisition/modules/journalctl/journalctl_test.go

@@ -21,6 +21,15 @@ import (
 	"github.com/crowdsecurity/crowdsec/pkg/pipeline"
 )
 
+// readTimeout is only ever reached when a line never arrives: it's generous on
+// purpose, so that a slow machine doesn't turn into a test failure.
+const readTimeout = 10 * time.Second
+
+// quietPeriod is how long we wait to confirm that no extra line shows up. Kept
+// short on purpose: a line slower than this is missed, which is better than
+// failing at random on a loaded machine.
+const quietPeriod = 100 * time.Millisecond
+
 func TestConfigureDSN(t *testing.T) {
 	cstest.SkipOnWindows(t)
 
@@ -147,33 +156,53 @@ journalctl_filter:
 	}
 	for idx, ts := range tests {
 		t.Run(strconv.Itoa(idx), func(t *testing.T) {
-			ctx, cancel := context.WithCancel(ctx)
+			// not shadowing ctx: the zombie check at the end needs a context that
+			// is still alive after the stream has been stopped
+			streamCtx, cancel := context.WithCancel(ctx)
 			out := make(chan pipeline.Event)
 			j := Source{}
 
 			logger, _ := logtest.NewNullLogger()
 
-			err := j.Configure(ctx, []byte(ts.config), logrus.NewEntry(logger), metrics.AcquisitionMetricsLevelNone)
+			err := j.Configure(streamCtx, []byte(ts.config), logrus.NewEntry(logger), metrics.AcquisitionMetricsLevelNone)
 			require.NoError(t, err)
 
-			gotLines := 0
-			var wg sync.WaitGroup
+			var (
+				wg         sync.WaitGroup
+				extraLines int
+			)
 
 			if ts.wantLines != 0 {
 				wg.Go(func() {
+					// whatever happens below, the stream has to be stopped or
+					// j.Stream() would never return
+					defer cancel()
+
+					// consuming every line lets journalctl exit on its own
+					for i := range ts.wantLines {
+						select {
+						case <-out:
+						case <-time.After(readTimeout):
+							t.Errorf("timed out waiting for line %d/%d", i+1, ts.wantLines)
+							return
+						}
+					}
+
+					// keep draining for a moment: it catches extra lines, and it keeps a
+					// source that sends too many from deadlocking on the unbuffered
+					// channel instead of failing the test
 					for {
 						select {
 						case <-out:
-							gotLines++
-						case <-time.After(1 * time.Second):
-							cancel()
+							extraLines++
+						case <-time.After(quietPeriod):
 							return
 						}
 					}
 				})
 			}
 
-			err = j.Stream(ctx, out)
+			err = j.Stream(streamCtx, out)
 			cstest.RequireErrorContains(t, err, ts.wantErr)
 
 			if ts.wantErr != "" {
@@ -183,11 +212,13 @@ journalctl_filter:
 
 			if ts.wantLines != 0 {
 				wg.Wait()
-				assert.Equal(t, ts.wantLines, gotLines)
+				assert.Zero(t, extraLines, "source emitted more lines than expected")
 			}
 
 			cancel()
 
+			// streamCtx is canceled by now, and running pgrep with it would make
+			// it return nothing and the check below always pass
 			output, _ := exec.CommandContext(ctx, "pgrep", "-x", "journalctl").CombinedOutput()
 			assert.Empty(t, output, "zombie journalctl process detected!")
 		})

pkg/acquisition/modules/kafka/kafka_test.go

@@ -18,6 +18,39 @@ import (
 	"github.com/crowdsecurity/crowdsec/pkg/pipeline"
 )
 
+// readTimeout is only ever reached when a message never arrives: it's generous on
+// purpose, so that a slow broker doesn't turn into a test failure.
+const readTimeout = 10 * time.Second
+
+// quietPeriod is how long we wait to confirm that no extra message shows up. Kept
+// short on purpose: an unexpected message slower than this is missed, which is
+// better than failing at random on a loaded machine.
+const quietPeriod = 100 * time.Millisecond
+
+// readEvents reads exactly n events, failing if they don't all show up in time.
+func readEvents(t *testing.T, out chan pipeline.Event, n int) {
+	t.Helper()
+
+	for i := range n {
+		select {
+		case <-out:
+		case <-time.After(readTimeout):
+			t.Fatalf("timed out waiting for message %d/%d", i+1, n)
+		}
+	}
+}
+
+// requireNoMoreEvents checks that the datasource doesn't emit more than expected.
+func requireNoMoreEvents(t *testing.T, out chan pipeline.Event) {
+	t.Helper()
+
+	select {
+	case evt := <-out:
+		t.Fatalf("unexpected extra message: %q", evt.Line.Raw)
+	case <-time.After(quietPeriod):
+	}
+}
+
 func writeToKafka(ctx context.Context, w *kafka.Writer, logs []string) {
 	for idx, log := range logs {
 		err := w.WriteMessages(ctx, kafka.Message{
@@ -118,20 +151,11 @@ topic: crowdsecplaintext`), subLogger, metrics.AcquisitionMetricsLevelNone)
 			err = k.StreamingAcquisition(ctx, out, &tomb)
 			cstest.AssertErrorContains(t, err, ts.expectedErr)
 
-			actualLines := 0
-
 			go writeToKafka(ctx, w, ts.logs)
-		READLOOP:
-			for {
-				select {
-				case <-out:
-					actualLines++
-				case <-time.After(2 * time.Second):
-					break READLOOP
-				}
-			}
 
-			require.Equal(t, ts.expectedLines, actualLines)
+			readEvents(t, out, ts.expectedLines)
+			requireNoMoreEvents(t, out)
+
 			tomb.Kill(nil)
 			err = tomb.Wait()
 			require.NoError(t, err)
@@ -196,20 +220,11 @@ tls:
 			err = k.StreamingAcquisition(ctx, out, &tomb)
 			cstest.AssertErrorContains(t, err, ts.expectedErr)
 
-			actualLines := 0
-
 			go writeToKafka(ctx, w2, ts.logs)
-		READLOOP:
-			for {
-				select {
-				case <-out:
-					actualLines++
-				case <-time.After(2 * time.Second):
-					break READLOOP
-				}
-			}
 
-			require.Equal(t, ts.expectedLines, actualLines)
+			readEvents(t, out, ts.expectedLines)
+			requireNoMoreEvents(t, out)
+
 			tomb.Kill(nil)
 			err = tomb.Wait()
 			require.NoError(t, err)

pkg/acquisition/modules/kubernetes/run_test.go

@@ -27,6 +27,11 @@ import (
 // asserting against the exact backoff duration.
 const retryBackoffMidpoint = 500 * time.Millisecond
 
+// waitTimeout is only ever reached when what we're waiting for never happens:
+// it's generous on purpose, so that a loaded machine doesn't turn into a test
+// failure.
+const waitTimeout = 10 * time.Second
+
 func testLogger() *log.Entry {
 	return log.WithField("type", ModuleName)
 }
@@ -183,7 +188,7 @@ func TestTailPod_StartsWorkerAndStopPodCancelsIt(t *testing.T) {
 	select {
 	case evt := <-out:
 		assert.Equal(t, "hello", evt.Line.Raw)
-	case <-time.After(time.Second):
+	case <-time.After(waitTimeout):
 		t.Fatal("timed out waiting for a log line from the tailed pod")
 	}
 
@@ -197,7 +202,7 @@ func TestTailPod_StartsWorkerAndStopPodCancelsIt(t *testing.T) {
 
 	select {
 	case <-done:
-	case <-time.After(time.Second):
+	case <-time.After(waitTimeout):
 		t.Fatal("stopPod did not stop the worker goroutine in time")
 	}