From 77c3a8a61e7e18adbdbbc759f983e3f34adf7b0f Mon Sep 17 00:00:00 2001
From: Lucas Manuel Rodriguez <lucarodriguez@gmail.com>
Date: Wed, 19 Jan 2022 16:17:00 -0300
Subject: [PATCH] Fix flaky `TestPolicyWebhooks` (#3777)

* Fix flaky TestPolicyWebhooks

* Run test redis cleanup before running tests
---
 server/datastore/redis/redistest/redistest.go | 30 ++++++++++++-------
 server/service/service_osquery_test.go        |  6 ++--
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/server/datastore/redis/redistest/redistest.go b/server/datastore/redis/redistest/redistest.go
index abd2fbe5a3..82d7ec0cf5 100644
--- a/server/datastore/redis/redistest/redistest.go
+++ b/server/datastore/redis/redistest/redistest.go
@@ -75,20 +75,28 @@ func SetupRedis(tb testing.TB, cleanupKeyPrefix string, cluster, redir, readRepl
 	_, err = conn.Do("PING")
 	require.Nil(tb, err)
 
+	// We run a cleanup before running the tests in case a previous
+	// run failed to run the cleanup (e.g. Ctrl+C while running tests).
+	cleanup(tb, pool, cleanupKeyPrefix)
+
 	tb.Cleanup(func() {
-		keys, err := redis.ScanKeys(pool, cleanupKeyPrefix+"*", 1000)
-		require.NoError(tb, err)
-		for _, k := range keys {
-			func() {
-				conn := pool.Get()
-				defer conn.Close()
-				if _, err := conn.Do("DEL", k); err != nil {
-					require.NoError(tb, err)
-				}
-			}()
-		}
+		cleanup(tb, pool, cleanupKeyPrefix)
 		pool.Close()
 	})
 
 	return pool
 }
+
+func cleanup(tb testing.TB, pool fleet.RedisPool, cleanupKeyPrefix string) {
+	keys, err := redis.ScanKeys(pool, cleanupKeyPrefix+"*", 1000)
+	require.NoError(tb, err)
+	for _, k := range keys {
+		func() {
+			conn := pool.Get()
+			defer conn.Close()
+			if _, err := conn.Do("DEL", k); err != nil {
+				require.NoError(tb, err)
+			}
+		}()
+	}
+}
diff --git a/server/service/service_osquery_test.go b/server/service/service_osquery_test.go
index 283eb0c16a..41a36d3c50 100644
--- a/server/service/service_osquery_test.go
+++ b/server/service/service_osquery_test.go
@@ -2303,7 +2303,7 @@ func TestPolicyWebhooks(t *testing.T) {
 			}},
 		})
 		return err == nil
-	}, 5*time.Second, 250*time.Millisecond)
+	}, 1*time.Minute, 250*time.Millisecond)
 	require.NoError(t, err)
 
 	noPolicyResults := func(queries map[string]string) {
@@ -2368,7 +2368,7 @@ func TestPolicyWebhooks(t *testing.T) {
 			3: {},
 		})
 		return err == nil
-	}, 5*time.Second, 250*time.Millisecond)
+	}, 1*time.Minute, 250*time.Millisecond)
 	require.NoError(t, err)
 
 	// Simulate webhook trigger by removing the hosts.
@@ -2408,7 +2408,7 @@ func TestPolicyWebhooks(t *testing.T) {
 			3: {},
 		})
 		return err == nil
-	}, 5*time.Second, 250*time.Millisecond)
+	}, 1*time.Minute, 250*time.Millisecond)
 	require.NoError(t, err)
 }