remove sleep

Author: aphralG

test/integration/grpc_management_plane_api_test.go

@@ -212,8 +212,6 @@ func TestGrpc_Reconnection(t *testing.T) {
 	require.NoError(t, err)
 	mockManagementPlaneAPIAddress = net.JoinHostPort(ipAddress, ports["9093/tcp"][0].HostPort)
 
-	time.Sleep(5 * time.Second)
-
 	currentID := verifyConnection(t, 2)
 	assert.Equal(t, originalID, currentID)
 }
@@ -534,15 +532,25 @@ func verifyConnection(t *testing.T, instancesLength int) string {
 
 	client := resty.New()
 	client.SetRetryCount(retryCount).SetRetryWaitTime(retryWaitTime).SetRetryMaxWaitTime(retryMaxWaitTime)
+	connectionRequest := mpi.CreateConnectionRequest{}
+	client.AddRetryCondition(
+		func(r *resty.Response, err error) bool {
+			responseData := r.Body()
+			assert.True(t, json.Valid(responseData))
 
+			pb := protojson.UnmarshalOptions{DiscardUnknown: true}
+			unmarshalErr := pb.Unmarshal(responseData, &connectionRequest)
+
+			return r.StatusCode() == http.StatusNotFound || unmarshalErr == nil
+		},
+	)
 	url := fmt.Sprintf("http://%s/api/v1/connection", mockManagementPlaneAPIAddress)
+	t.Logf("Connecting to %s", url)
 	resp, err := client.R().EnableTrace().Get(url)
 
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, resp.StatusCode())
 
-	connectionRequest := mpi.CreateConnectionRequest{}
-
 	responseData := resp.Body()
 	t.Logf("Response: %s", string(responseData))
 	assert.True(t, json.Valid(responseData))