From f801244453955672f5038636a25e3e9b8eb38741 Mon Sep 17 00:00:00 2001
From: William Beuil <william.beuil@gmail.com>
Date: Tue, 1 Feb 2022 11:16:11 +0100
Subject: [PATCH] chore: add retries on flaky acc tests

---
 pkg/resource/aws/aws_eip_test.go              | 11 +++++++++++
 pkg/resource/aws/aws_internet_gateway_test.go | 11 +++++++++++
 pkg/resource/aws/aws_network_acl_test.go      | 12 ++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/pkg/resource/aws/aws_eip_test.go b/pkg/resource/aws/aws_eip_test.go
index 1f16ad5fc..01293acbc 100644
--- a/pkg/resource/aws/aws_eip_test.go
+++ b/pkg/resource/aws/aws_eip_test.go
@@ -22,6 +22,17 @@ func TestAcc_Aws_Eip(t *testing.T) {
 				Env: map[string]string{
 					"AWS_REGION": "us-east-1",
 				},
+				// New resources are not visible immediately through AWS API after an apply operation
+				// (e.g. error attaching EC2 Internet Gateway)
+				// Logic below retries driftctl scan using a back-off strategy of retrying 'n' times
+				// and doubling the amount of time waited after each one.
+				ShouldRetry: func(result *test.ScanResult, retryDuration time.Duration, retryCount uint8) bool {
+					if result.IsSync() || retryDuration > 10*time.Minute {
+						return false
+					}
+					time.Sleep((2 * time.Duration(retryCount)) * time.Minute)
+					return true
+				},
 				Check: func(result *test.ScanResult, stdout string, err error) {
 					if err != nil {
 						t.Fatal(err)
diff --git a/pkg/resource/aws/aws_internet_gateway_test.go b/pkg/resource/aws/aws_internet_gateway_test.go
index 1b97e1ed7..b7915ae76 100644
--- a/pkg/resource/aws/aws_internet_gateway_test.go
+++ b/pkg/resource/aws/aws_internet_gateway_test.go
@@ -22,6 +22,17 @@ func TestAcc_Aws_InternetGateway(t *testing.T) {
 				Env: map[string]string{
 					"AWS_REGION": "us-east-1",
 				},
+				// New resources are not visible immediately through AWS API after an apply operation
+				// (e.g. InvalidVpcID.NotFound)
+				// Logic below retries driftctl scan using a back-off strategy of retrying 'n' times
+				// and doubling the amount of time waited after each one.
+				ShouldRetry: func(result *test.ScanResult, retryDuration time.Duration, retryCount uint8) bool {
+					if result.IsSync() || retryDuration > 10*time.Minute {
+						return false
+					}
+					time.Sleep((2 * time.Duration(retryCount)) * time.Minute)
+					return true
+				},
 				Check: func(result *test.ScanResult, stdout string, err error) {
 					if err != nil {
 						t.Fatal(err)
diff --git a/pkg/resource/aws/aws_network_acl_test.go b/pkg/resource/aws/aws_network_acl_test.go
index 8a23e3d0d..1b81e3f72 100644
--- a/pkg/resource/aws/aws_network_acl_test.go
+++ b/pkg/resource/aws/aws_network_acl_test.go
@@ -2,6 +2,7 @@ package aws_test
 
 import (
 	"testing"
+	"time"
 
 	"github.com/snyk/driftctl/test"
 	"github.com/snyk/driftctl/test/acceptance"
@@ -41,6 +42,17 @@ func TestAcc_Aws_NetworkAcl_NonDeep(t *testing.T) {
 				Env: map[string]string{
 					"AWS_REGION": "us-east-1",
 				},
+				// New resources are not visible immediately through AWS API after an apply operation
+				// (e.g. InvalidNetworkAclID.NotFound)
+				// Logic below retries driftctl scan using a back-off strategy of retrying 'n' times
+				// and doubling the amount of time waited after each one.
+				ShouldRetry: func(result *test.ScanResult, retryDuration time.Duration, retryCount uint8) bool {
+					if result.IsSync() || retryDuration > 10*time.Minute {
+						return false
+					}
+					time.Sleep((2 * time.Duration(retryCount)) * time.Minute)
+					return true
+				},
 				Check: func(result *test.ScanResult, stdout string, err error) {
 					if err != nil {
 						t.Fatal(err)