Double up the athena tests that hit AWS

lib/events/athena/integration_test.go

@@ -37,9 +37,20 @@ import (
 )
 
 func TestIntegrationAthenaSearchSessionEventsBySessionID(t *testing.T) {
+	t.Run("sns", func(t *testing.T) {
+		const bypassSNSFalse = false
+		testIntegrationAthenaSearchSessionEventsBySessionID(t, bypassSNSFalse)
+	})
+	t.Run("sqs", func(t *testing.T) {
+		const bypassSNSTrue = true
+		testIntegrationAthenaSearchSessionEventsBySessionID(t, bypassSNSTrue)
+	})
+}
+
+func testIntegrationAthenaSearchSessionEventsBySessionID(t *testing.T, bypassSNS bool) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
-	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{})
+	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{BypassSNS: bypassSNS})
 	auditLogger := &EventuallyConsistentAuditLogger{
 		Inner: ac.log,
 		// Additional 5s is used to compensate for uploading parquet on s3.
@@ -55,9 +66,20 @@ func TestIntegrationAthenaSearchSessionEventsBySessionID(t *testing.T) {
 }
 
 func TestIntegrationAthenaSessionEventsCRUD(t *testing.T) {
+	t.Run("sns", func(t *testing.T) {
+		const bypassSNSFalse = false
+		testIntegrationAthenaSessionEventsCRUD(t, bypassSNSFalse)
+	})
+	t.Run("sqs", func(t *testing.T) {
+		const bypassSNSTrue = true
+		testIntegrationAthenaSessionEventsCRUD(t, bypassSNSTrue)
+	})
+}
+
+func testIntegrationAthenaSessionEventsCRUD(t *testing.T, bypassSNS bool) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
-	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{})
+	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{BypassSNS: bypassSNS})
 	auditLogger := &EventuallyConsistentAuditLogger{
 		Inner: ac.log,
 		// Additional 5s is used to compensate for uploading parquet on s3.
@@ -72,9 +94,20 @@ func TestIntegrationAthenaSessionEventsCRUD(t *testing.T) {
 }
 
 func TestIntegrationAthenaEventPagination(t *testing.T) {
+	t.Run("sns", func(t *testing.T) {
+		const bypassSNSFalse = false
+		testIntegrationAthenaEventPagination(t, bypassSNSFalse)
+	})
+	t.Run("sqs", func(t *testing.T) {
+		const bypassSNSTrue = true
+		testIntegrationAthenaEventPagination(t, bypassSNSTrue)
+	})
+}
+
+func testIntegrationAthenaEventPagination(t *testing.T, bypassSNS bool) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
-	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{})
+	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{BypassSNS: bypassSNS})
 	auditLogger := &EventuallyConsistentAuditLogger{
 		Inner: ac.log,
 		// Additional 5s is used to compensate for uploading parquet on s3.
@@ -89,10 +122,24 @@ func TestIntegrationAthenaEventPagination(t *testing.T) {
 }
 
 func TestIntegrationAthenaLargeEvents(t *testing.T) {
+	t.Run("sns", func(t *testing.T) {
+		const bypassSNSFalse = false
+		testIntegrationAthenaLargeEvents(t, bypassSNSFalse)
+	})
+	t.Run("sqs", func(t *testing.T) {
+		const bypassSNSTrue = true
+		testIntegrationAthenaLargeEvents(t, bypassSNSTrue)
+	})
+}
+
+func testIntegrationAthenaLargeEvents(t *testing.T, bypassSNS bool) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{MaxBatchSize: 1})
+	ac := SetupAthenaContext(t, ctx, AthenaContextConfig{
+		MaxBatchSize: 1,
+		BypassSNS:    bypassSNS,
+	})
 	in := &apievents.SessionStart{
 		Metadata: apievents.Metadata{
 			Index: 2,

lib/events/athena/test.go

@@ -106,6 +106,7 @@ func (a *AthenaContext) GetLog() *Log {
 // AthenaContextConfig is optional config to override defaults in athena context.
 type AthenaContextConfig struct {
 	MaxBatchSize int
+	BypassSNS    bool
 }
 
 type InfraOutputs struct {
@@ -169,13 +170,6 @@ func SetupAthenaContext(t *testing.T, ctx context.Context, cfg AthenaContextConf
 	if ok, _ := strconv.ParseBool(testEnabled); !ok {
 		t.Skip("Skipping AWS-dependent test suite.")
 	}
-	var bypassSNS bool
-	if s := os.Getenv(teleport.AWSRunTests + "_ATHENA_BYPASS_SNS"); s != "" {
-		if ok, _ := strconv.ParseBool(s); ok {
-			t.Log("bypassing SNS for Athena audit log events")
-			bypassSNS = true
-		}
-	}
 
 	testID := fmt.Sprintf("auditlogs-integrationtests-%v", uuid.New().String())
 
@@ -208,7 +202,7 @@ func SetupAthenaContext(t *testing.T, ctx context.Context, cfg AthenaContextConf
 	}
 
 	topicARN := infraOut.TopicARN
-	if bypassSNS {
+	if cfg.BypassSNS {
 		topicARN = topicARNBypass
 	}
 	log, err := New(ctx, Config{