fix: try to deflake several flaky tests

app/app_test.go

@@ -245,25 +245,14 @@ func TestAppIntegration(t *testing.T) {
 	err = startstop.Stop(graph.Objects(), nil)
 	assert.NoError(t, err)
 
-	// Wait for span to be sent.
-	deadline := time.After(time.Second)
-	for {
-		if out.Len() > 62 {
-			break
-		}
-		select {
-		case <-deadline:
-			t.Error("timed out waiting for output")
-			return
-		case <-time.After(time.Millisecond):
-		}
-	}
+	assert.Eventually(t, func() bool {
+		return out.Len() > 62
+	}, 5*time.Second, 2*time.Millisecond)
 	assert.Equal(t, `{"data":{"foo":"bar","meta.refinery.original_sample_rate":1,"trace.trace_id":"1"},"dataset":"dataset"}`+"\n", out.String())
 }
 
 func TestAppIntegrationWithNonLegacyKey(t *testing.T) {
-	// This is failing in Parallel, so disable it for now.
-	// t.Parallel()
+	t.Parallel()
 
 	var out bytes.Buffer
 	a, graph := newStartedApp(t, &transmission.WriterSender{W: &out}, 10500, nil, false)
@@ -288,18 +277,9 @@ func TestAppIntegrationWithNonLegacyKey(t *testing.T) {
 	assert.NoError(t, err)
 
 	// Wait for span to be sent.
-	deadline := time.After(2 * time.Second)
-	for {
-		if out.Len() > 62 {
-			break
-		}
-		select {
-		case <-deadline:
-			t.Error("timed out waiting for output")
-			return
-		case <-time.After(time.Millisecond):
-		}
-	}
+	assert.Eventually(t, func() bool {
+		return out.Len() > 62
+	}, 5*time.Second, 2*time.Millisecond)
 	assert.Equal(t, `{"data":{"foo":"bar","meta.refinery.original_sample_rate":1,"trace.trace_id":"1"},"dataset":"dataset"}`+"\n", out.String())
 }
 
@@ -452,19 +432,9 @@ func TestHostMetadataSpanAdditions(t *testing.T) {
 	err = startstop.Stop(graph.Objects(), nil)
 	assert.NoError(t, err)
 
-	// Wait for span to be sent.
-	deadline := time.After(time.Second)
-	for {
-		if out.Len() > 62 {
-			break
-		}
-		select {
-		case <-deadline:
-			t.Error("timed out waiting for output")
-			return
-		case <-time.After(time.Millisecond):
-		}
-	}
+	assert.Eventually(t, func() bool {
+		return out.Len() > 62
+	}, 5*time.Second, 2*time.Millisecond)
 
 	expectedSpan := `{"data":{"foo":"bar","meta.refinery.local_hostname":"%s","meta.refinery.original_sample_rate":1,"trace.trace_id":"1"},"dataset":"dataset"}` + "\n"
 	assert.Equal(t, fmt.Sprintf(expectedSpan, hostname), out.String())

collect/collect_test.go

@@ -152,7 +152,7 @@ func TestOriginalSampleRateIsNotedInMetaField(t *testing.T) {
 
 	// Spin until a sample gets triggered
 	sendAttemptCount := 0
-	for getEventsLength(transmission) < 1 || sendAttemptCount > 10 {
+	for getEventsLength(transmission) < 1 {
 		sendAttemptCount++
 		span := &types.Span{
 			TraceID: fmt.Sprintf("trace-%v", sendAttemptCount),
@@ -164,11 +164,11 @@ func TestOriginalSampleRateIsNotedInMetaField(t *testing.T) {
 			},
 		}
 		coll.AddSpan(span)
-		time.Sleep(conf.SendTickerVal * 2)
+		time.Sleep(conf.SendTickerVal * 5)
 	}
 
 	transmission.Mux.RLock()
-	assert.Greater(t, len(transmission.Events), 0, "should be some events transmitted")
+	assert.Greater(t, len(transmission.Events), 0, "should be at least one event transmitted")
 	assert.Equal(t, uint(50), transmission.Events[0].Data["meta.refinery.original_sample_rate"],
 		"metadata should be populated with original sample rate")
 	transmission.Mux.RUnlock()
@@ -245,8 +245,13 @@ func TestTransmittedSpansShouldHaveASampleRateOfAtLeastOne(t *testing.T) {
 
 	time.Sleep(conf.SendTickerVal * 2)
 
+	assert.Eventually(t, func() bool {
+		transmission.Mux.RLock()
+		defer transmission.Mux.RUnlock()
+		return len(transmission.Events) > 0
+	}, 2*time.Second, conf.SendTickerVal*2)
+
 	transmission.Mux.RLock()
-	assert.Equal(t, 1, len(transmission.Events), "should be some events transmitted")
 	assert.Equal(t, uint(1), transmission.Events[0].SampleRate,
 		"SampleRate should be reset to one after starting at zero")
 	transmission.Mux.RUnlock()

internal/peer/peers_test.go

@@ -70,8 +70,10 @@ func TestPeerShutdown(t *testing.T) {
 	assert.True(t, strings.HasSuffix(peers[0], "8081"))
 
 	close(done)
-	time.Sleep(500 * time.Millisecond)
-	peers, err = peer.GetPeers()
-	assert.NoError(t, err)
-	assert.Equal(t, 0, len(peers))
+
+	assert.Eventually(t, func() bool {
+		peers, err = peer.GetPeers()
+		assert.NoError(t, err)
+		return len(peers) == 0
+	}, 5*time.Second, 200*time.Millisecond)
 }