From 58a492d56d7c7efd84158a4eba8220bcf776bbfa Mon Sep 17 00:00:00 2001
From: David Colburn <xero73@gmail.com>
Date: Thu, 8 Aug 2024 17:58:54 -0400
Subject: [PATCH] fix test races for stream updates

---
 test/edge.go        | 10 +++++++---
 test/integration.go | 32 ++++++++++++++++++++++++--------
 test/stream.go      |  8 --------
 3 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/test/edge.go b/test/edge.go
index 4677d252..8ad96bc5 100644
--- a/test/edge.go
+++ b/test/edge.go
@@ -150,10 +150,14 @@ func (r *Runner) testRtmpFailure(t *testing.T) {
 		time.Sleep(time.Second * 5)
 		info = r.getUpdate(t, info.EgressId)
 		if info.Status == livekit.EgressStatus_EGRESS_ACTIVE {
-			r.checkUpdate(t, info.EgressId, livekit.EgressStatus_EGRESS_FAILED)
-		} else {
-			require.Equal(t, livekit.EgressStatus_EGRESS_FAILED, info.Status)
+			info = r.getUpdate(t, info.EgressId)
+		}
+		if info.Status == livekit.EgressStatus_EGRESS_ACTIVE {
+			require.Equal(t, livekit.StreamInfo_FAILED, info.StreamResults[0].Status)
+			info = r.getUpdate(t, info.EgressId)
 		}
+		require.Equal(t, livekit.StreamInfo_FAILED, info.StreamResults[0].Status)
+		require.Equal(t, livekit.EgressStatus_EGRESS_FAILED, info.Status)
 	})
 }
 
diff --git a/test/integration.go b/test/integration.go
index 0285b8f3..ae99fc01 100644
--- a/test/integration.go
+++ b/test/integration.go
@@ -146,16 +146,32 @@ func (r *Runner) checkUpdate(t *testing.T, egressID string, status livekit.Egres
 }
 
 func (r *Runner) checkStreamUpdate(t *testing.T, egressID string, expected map[string]livekit.StreamInfo_Status) {
-	info := r.getUpdate(t, egressID)
+	for {
+		info := r.getUpdate(t, egressID)
+		require.Equal(t, len(expected), len(info.StreamResults))
+
+		failureStillActive := false
+		for _, s := range info.StreamResults {
+			require.Equal(t, s.Status == livekit.StreamInfo_FAILED, s.Error != "")
+
+			var e livekit.StreamInfo_Status
+			if strings.HasSuffix(s.Url, ".contribute.live-video.net/app/{f...1}") {
+				e = expected[badRtmpUrl1Redacted]
+			} else {
+				e = expected[s.Url]
+			}
+			if e == livekit.StreamInfo_FAILED && s.Status == livekit.StreamInfo_ACTIVE {
+				// expecting another update
+				failureStillActive = true
+				continue
+			}
 
-	require.Equal(t, len(expected), len(info.StreamResults))
-	for _, s := range info.StreamResults {
-		if strings.HasSuffix(s.Url, ".contribute.live-video.net/app/{f...1}") {
-			require.Equal(t, expected[badRtmpUrl1Redacted], s.Status)
-		} else {
-			require.Equal(t, expected[s.Url], s.Status)
+			require.Equal(t, e, s.Status)
+		}
+
+		if !failureStillActive {
+			return
 		}
-		require.Equal(t, s.Status == livekit.StreamInfo_FAILED, s.Error != "")
 	}
 }
 
diff --git a/test/stream.go b/test/stream.go
index be2385ef..da09e640 100644
--- a/test/stream.go
+++ b/test/stream.go
@@ -92,14 +92,6 @@ func (r *Runner) runStreamTest(t *testing.T, req *rpc.StartEgressRequest, test *
 
 	// verify
 	r.verifyStreams(t, p, urls[0][2], urls[2][2])
-	if test.outputType != types.OutputTypeRTMP {
-		r.checkStreamUpdate(t, egressID, map[string]livekit.StreamInfo_Status{
-			urls[0][1]: livekit.StreamInfo_ACTIVE,
-			urls[1][1]: livekit.StreamInfo_FAILED,
-			urls[2][1]: livekit.StreamInfo_ACTIVE,
-			urls[3][1]: livekit.StreamInfo_ACTIVE,
-		})
-	}
 	r.checkStreamUpdate(t, egressID, map[string]livekit.StreamInfo_Status{
 		urls[0][1]: livekit.StreamInfo_ACTIVE,
 		urls[1][1]: livekit.StreamInfo_FAILED,