From 73da6d15f94a1004d6be9296b60d43f24c713ce7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Tyczy=C5=84ski?= <wojtekt@google.com>
Date: Tue, 5 Apr 2022 20:40:38 +0200
Subject: [PATCH] Fix TestPriorityLevelIsolation concurrency issue

---
 .../apiserver/flowcontrol/concurrency_test.go    | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/test/integration/apiserver/flowcontrol/concurrency_test.go b/test/integration/apiserver/flowcontrol/concurrency_test.go
index 7c5a11be9d4..d5f314af22b 100644
--- a/test/integration/apiserver/flowcontrol/concurrency_test.go
+++ b/test/integration/apiserver/flowcontrol/concurrency_test.go
@@ -22,6 +22,7 @@ import (
 	"io"
 	"net/http/httptest"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -104,22 +105,28 @@ func TestPriorityLevelIsolation(t *testing.T) {
 	}
 
 	stopCh := make(chan struct{})
-	defer close(stopCh)
+	wg := sync.WaitGroup{}
+	defer func() {
+		close(stopCh)
+		wg.Wait()
+	}()
 
 	// "elephant"
+	wg.Add(concurrencyShares + queueLength)
 	streamRequests(concurrencyShares+queueLength, func() {
 		_, err := noxu1Client.CoreV1().Namespaces().List(context.Background(), metav1.ListOptions{})
 		if err != nil {
 			t.Error(err)
 		}
-	}, stopCh)
+	}, &wg, stopCh)
 	// "mouse"
+	wg.Add(3)
 	streamRequests(3, func() {
 		_, err := noxu2Client.CoreV1().Namespaces().List(context.Background(), metav1.ListOptions{})
 		if err != nil {
 			t.Error(err)
 		}
-	}, stopCh)
+	}, &wg, stopCh)
 
 	time.Sleep(time.Second * 10) // running in background for a while
 
@@ -312,9 +319,10 @@ func createPriorityLevelAndBindingFlowSchemaForUser(c clientset.Interface, usern
 	})
 }
 
-func streamRequests(parallel int, request func(), stopCh <-chan struct{}) {
+func streamRequests(parallel int, request func(), wg *sync.WaitGroup, stopCh <-chan struct{}) {
 	for i := 0; i < parallel; i++ {
 		go func() {
+			defer wg.Done()
 			for {
 				select {
 				case <-stopCh: