fix flaky test TestSubscriptions (#2779)

* fix flaky test TestSubscriptions * update other copy of the test
99designs · Sep 5, 2023 · bee47dc · bee47dc
1 parent a41f4da
commit bee47dc
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/codegen/testserver/followschema/subscription_test.go b/codegen/testserver/followschema/subscription_test.go
@@ -130,6 +130,9 @@ func TestSubscriptions(t *testing.T) {
 	})
 
 	t.Run("will parse init payload", func(t *testing.T) {
+		runtime.GC() // ensure no go-routines left from preceding tests
+		initialGoroutineCount := runtime.NumGoroutine()
+
 		sub := c.WebsocketWithPayload(`subscription { initPayload }`, map[string]interface{}{
 			"Authorization": "Bearer of the curse",
 			"number":        32,
@@ -155,6 +158,14 @@ func TestSubscriptions(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "strings = []interface {}{\"hello\", \"world\"}", msg.resp.InitPayload)
 		sub.Close()
+
+		// need a little bit of time for goroutines to settle
+		start := time.Now()
+		for time.Since(start).Seconds() < 2 && initialGoroutineCount != runtime.NumGoroutine() {
+			time.Sleep(5 * time.Millisecond)
+		}
+
+		require.Equal(t, initialGoroutineCount, runtime.NumGoroutine())
 	})
 
 	t.Run("websocket gets errors", func(t *testing.T) {

diff --git a/codegen/testserver/singlefile/subscription_test.go b/codegen/testserver/singlefile/subscription_test.go
@@ -130,6 +130,9 @@ func TestSubscriptions(t *testing.T) {
 	})
 
 	t.Run("will parse init payload", func(t *testing.T) {
+		runtime.GC() // ensure no go-routines left from preceding tests
+		initialGoroutineCount := runtime.NumGoroutine()
+
 		sub := c.WebsocketWithPayload(`subscription { initPayload }`, map[string]interface{}{
 			"Authorization": "Bearer of the curse",
 			"number":        32,
@@ -155,6 +158,14 @@ func TestSubscriptions(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "strings = []interface {}{\"hello\", \"world\"}", msg.resp.InitPayload)
 		sub.Close()
+
+		// need a little bit of time for goroutines to settle
+		start := time.Now()
+		for time.Since(start).Seconds() < 2 && initialGoroutineCount != runtime.NumGoroutine() {
+			time.Sleep(5 * time.Millisecond)
+		}
+
+		require.Equal(t, initialGoroutineCount, runtime.NumGoroutine())
 	})
 
 	t.Run("websocket gets errors", func(t *testing.T) {