diff --git a/securitycenter/management_api/event_threat_detection/event_threat_detection_custom_module_test.go b/securitycenter/management_api/event_threat_detection/event_threat_detection_custom_module_test.go index fa2cfb85aa..e5d245d9a4 100644 --- a/securitycenter/management_api/event_threat_detection/event_threat_detection_custom_module_test.go +++ b/securitycenter/management_api/event_threat_detection/event_threat_detection_custom_module_test.go @@ -22,6 +22,7 @@ import ( "math/rand" "os" "strings" + "sync" "testing" "time" @@ -33,6 +34,7 @@ import ( var orgID = "" var createdCustomModuleID = "" +var mu sync.Mutex func TestMain(m *testing.M) { orgID = os.Getenv("GCLOUD_ORGANIZATION") @@ -54,14 +56,67 @@ func TestMain(m *testing.M) { os.Exit(code) } -// extractCustomModuleID extracts the custom module ID from the full name -func extractCustomModuleID(customModuleFullName string) string { - trimmedFullName := strings.TrimSpace(customModuleFullName) - parts := strings.Split(trimmedFullName, "/") - if len(parts) > 0 { - return parts[len(parts)-1] +func cleanupExistingCustomModules(orgID string) error { + ctx := context.Background() + client, err := securitycentermanagement.NewClient(ctx) + if err != nil { + return fmt.Errorf("securitycentermanagement.NewClient: %w", err) } - return "" + defer client.Close() + + parent := fmt.Sprintf("organizations/%s/locations/global", orgID) + + // List all existing custom modules + req := &securitycentermanagementpb.ListEventThreatDetectionCustomModulesRequest{ + Parent: parent, + } + + it := client.ListEventThreatDetectionCustomModules(ctx, req) + for { + module, err := it.Next() + + if err == iterator.Done { + break + } + + if err != nil { + return fmt.Errorf("failed to list CustomModules: %w", err) + } + + // Check if the custom module name starts with 'go_sample_etd_custom' + if strings.HasPrefix(module.DisplayName, "go_sample_etd_custom") { + + customModuleID := extractCustomModuleID(module.Name) + // Delete the custom module + err := cleanupCustomModule(customModuleID) + if err != nil { + return fmt.Errorf("failed to delete existing CustomModule: %w", err) + } + fmt.Printf("Deleted existing CustomModule: %s\n", module.Name) + } + } + + return nil +} + +func cleanupCustomModule(customModuleID string) error { + + ctx := context.Background() + client, err := securitycentermanagement.NewClient(ctx) + if err != nil { + return fmt.Errorf("securitycentermanagement.NewClient: %w", err) + } + defer client.Close() + + req := &securitycentermanagementpb.DeleteEventThreatDetectionCustomModuleRequest{ + Name: fmt.Sprintf("organizations/%s/locations/global/eventThreatDetectionCustomModules/%s", orgID, customModuleID), + } + + if err := client.DeleteEventThreatDetectionCustomModule(ctx, req); err != nil { + return fmt.Errorf("failed to delete CustomModule: %w", err) + } + + return nil } // addCustomModule creates a custom module for testing purposes @@ -130,84 +185,24 @@ func addCustomModule() (string, error) { return createdCustomModuleID, nil } -func cleanupCustomModule(customModuleID string) error { - - ctx := context.Background() - client, err := securitycentermanagement.NewClient(ctx) - if err != nil { - return fmt.Errorf("securitycentermanagement.NewClient: %w", err) - } - defer client.Close() - - req := &securitycentermanagementpb.DeleteEventThreatDetectionCustomModuleRequest{ - Name: fmt.Sprintf("organizations/%s/locations/global/eventThreatDetectionCustomModules/%s", orgID, customModuleID), - } - - if err := client.DeleteEventThreatDetectionCustomModule(ctx, req); err != nil { - return fmt.Errorf("failed to delete CustomModule: %w", err) - } - - return nil -} - -func cleanupExistingCustomModules(orgID string) error { - ctx := context.Background() - client, err := securitycentermanagement.NewClient(ctx) - if err != nil { - return fmt.Errorf("securitycentermanagement.NewClient: %w", err) - } - defer client.Close() - - parent := fmt.Sprintf("organizations/%s/locations/global", orgID) - - // List all existing custom modules - req := &securitycentermanagementpb.ListEventThreatDetectionCustomModulesRequest{ - Parent: parent, - } - - it := client.ListEventThreatDetectionCustomModules(ctx, req) - for { - module, err := it.Next() - - if err == iterator.Done { - break - } - - if err != nil { - return fmt.Errorf("failed to list CustomModules: %w", err) - } - - // Check if the custom module name starts with 'go_sample_etd_custom' - if strings.HasPrefix(module.DisplayName, "go_sample_etd_custom") { - - customModuleID := extractCustomModuleID(module.Name) - // Delete the custom module - err := cleanupCustomModule(customModuleID) - if err != nil { - return fmt.Errorf("failed to delete existing CustomModule: %w", err) - } - fmt.Printf("Deleted existing CustomModule: %s\n", module.Name) - } +// extractCustomModuleID extracts the custom module ID from the full name +func extractCustomModuleID(customModuleFullName string) string { + trimmedFullName := strings.TrimSpace(customModuleFullName) + parts := strings.Split(trimmedFullName, "/") + if len(parts) > 0 { + return parts[len(parts)-1] } - - return nil + return "" } // TestCreateEtdCustomModule verifies the Create functionality func TestCreateEtdCustomModule(t *testing.T) { var buf bytes.Buffer - _, err := addCustomModule() - - if err != nil { - t.Fatalf("Could not setup test environment: %v", err) - return - } - parent := fmt.Sprintf("organizations/%s/locations/global", orgID) // Call Create - err = createEventThreatDetectionCustomModule(&buf, parent) + err := createEventThreatDetectionCustomModule(&buf, parent) if err != nil { t.Fatalf("createCustomModule() had error: %v", err) @@ -225,10 +220,13 @@ func TestCreateEtdCustomModule(t *testing.T) { func TestGetEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestGetEtdCustomModule: %v", err) return } @@ -254,10 +252,13 @@ func TestGetEtdCustomModule(t *testing.T) { func TestUpdateEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestUpdateEtdCustomModule: %v", err) return } @@ -281,10 +282,13 @@ func TestUpdateEtdCustomModule(t *testing.T) { func TestDeleteEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestDeleteEtdCustomModule: %v", err) return } @@ -308,10 +312,13 @@ func TestDeleteEtdCustomModule(t *testing.T) { func TestListEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + _, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestListEtdCustomModule: %v", err) return } @@ -336,10 +343,13 @@ func TestListEtdCustomModule(t *testing.T) { func TestListEffectiveEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + _, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestListEffectiveEtdCustomModule: %v", err) return } @@ -364,10 +374,13 @@ func TestListEffectiveEtdCustomModule(t *testing.T) { func TestGetEffectiveEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestGetEffectiveEtdCustomModule: %v", err) return } @@ -393,10 +406,13 @@ func TestGetEffectiveEtdCustomModule(t *testing.T) { func TestListDescendantEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + _, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestListDescendantEtdCustomModule: %v", err) return } @@ -421,10 +437,13 @@ func TestListDescendantEtdCustomModule(t *testing.T) { func TestValidateEtdCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + _, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestValidateEtdCustomModule: %v", err) return } diff --git a/securitycenter/management_api/security_health_analytics_custom_module_test.go b/securitycenter/management_api/security_health_analytics_custom_module_test.go index 35d985cfaf..a292b7f9f0 100644 --- a/securitycenter/management_api/security_health_analytics_custom_module_test.go +++ b/securitycenter/management_api/security_health_analytics_custom_module_test.go @@ -22,6 +22,7 @@ import ( "math/rand" "os" "strings" + "sync" "testing" "time" @@ -33,6 +34,7 @@ import ( var orgID = "" var createdCustomModuleID = "" +var mu sync.Mutex func TestMain(m *testing.M) { orgID = os.Getenv("GCLOUD_ORGANIZATION") @@ -54,14 +56,67 @@ func TestMain(m *testing.M) { os.Exit(code) } -// extractCustomModuleID extracts the custom module ID from the full name -func extractCustomModuleID(customModuleFullName string) string { - trimmedFullName := strings.TrimSpace(customModuleFullName) - parts := strings.Split(trimmedFullName, "/") - if len(parts) > 0 { - return parts[len(parts)-1] +func cleanupExistingCustomModules(orgID string) error { + ctx := context.Background() + client, err := securitycentermanagement.NewClient(ctx) + if err != nil { + return fmt.Errorf("securitycentermanagement.NewClient: %w", err) } - return "" + defer client.Close() + + parent := fmt.Sprintf("organizations/%s/locations/global", orgID) + + // List all existing custom modules + req := &securitycentermanagementpb.ListSecurityHealthAnalyticsCustomModulesRequest{ + Parent: parent, + } + + it := client.ListSecurityHealthAnalyticsCustomModules(ctx, req) + for { + module, err := it.Next() + + if err == iterator.Done { + break + } + + if err != nil { + return fmt.Errorf("failed to list CustomModules: %w", err) + } + + // Check if the custom module name starts with 'go_sample_sha_custom' + if strings.HasPrefix(module.DisplayName, "go_sample_sha_custom") { + + customModuleID := extractCustomModuleID(module.Name) + // Delete the custom module + err := cleanupCustomModule(customModuleID) + if err != nil { + return fmt.Errorf("failed to delete existing CustomModule: %w", err) + } + fmt.Printf("Deleted existing CustomModule: %s\n", module.Name) + } + } + + return nil +} + +func cleanupCustomModule(customModuleID string) error { + + ctx := context.Background() + client, err := securitycentermanagement.NewClient(ctx) + if err != nil { + return fmt.Errorf("securitycentermanagement.NewClient: %w", err) + } + defer client.Close() + + req := &securitycentermanagementpb.DeleteSecurityHealthAnalyticsCustomModuleRequest{ + Name: fmt.Sprintf("organizations/%s/locations/global/securityHealthAnalyticsCustomModules/%s", orgID, customModuleID), + } + + if err := client.DeleteSecurityHealthAnalyticsCustomModule(ctx, req); err != nil { + return fmt.Errorf("failed to delete CustomModule: %w", err) + } + + return nil } // addCustomModule creates a custom module for testing purposes @@ -137,73 +192,23 @@ func addCustomModule() (string, error) { return createdCustomModuleID, nil } -func cleanupCustomModule(customModuleID string) error { - - ctx := context.Background() - client, err := securitycentermanagement.NewClient(ctx) - if err != nil { - return fmt.Errorf("securitycentermanagement.NewClient: %w", err) - } - defer client.Close() - - req := &securitycentermanagementpb.DeleteSecurityHealthAnalyticsCustomModuleRequest{ - Name: fmt.Sprintf("organizations/%s/locations/global/securityHealthAnalyticsCustomModules/%s", orgID, customModuleID), - } - - if err := client.DeleteSecurityHealthAnalyticsCustomModule(ctx, req); err != nil { - return fmt.Errorf("failed to delete CustomModule: %w", err) - } - - return nil -} - -func cleanupExistingCustomModules(orgID string) error { - ctx := context.Background() - client, err := securitycentermanagement.NewClient(ctx) - if err != nil { - return fmt.Errorf("securitycentermanagement.NewClient: %w", err) - } - defer client.Close() - - parent := fmt.Sprintf("organizations/%s/locations/global", orgID) - - // List all existing custom modules - req := &securitycentermanagementpb.ListSecurityHealthAnalyticsCustomModulesRequest{ - Parent: parent, - } - - it := client.ListSecurityHealthAnalyticsCustomModules(ctx, req) - for { - module, err := it.Next() - - if err == iterator.Done { - break - } - - if err != nil { - return fmt.Errorf("failed to list CustomModules: %w", err) - } - - // Check if the custom module name starts with 'go_sample_sha_custom' - if strings.HasPrefix(module.DisplayName, "go_sample_sha_custom") { - - customModuleID := extractCustomModuleID(module.Name) - // Delete the custom module - err := cleanupCustomModule(customModuleID) - if err != nil { - return fmt.Errorf("failed to delete existing CustomModule: %w", err) - } - fmt.Printf("Deleted existing CustomModule: %s\n", module.Name) - } +// extractCustomModuleID extracts the custom module ID from the full name +func extractCustomModuleID(customModuleFullName string) string { + trimmedFullName := strings.TrimSpace(customModuleFullName) + parts := strings.Split(trimmedFullName, "/") + if len(parts) > 0 { + return parts[len(parts)-1] } - - return nil + return "" } // TestDeleteCustomModule verifies the List functionality func TestDeleteCustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { @@ -227,8 +232,8 @@ func TestDeleteCustomModule(t *testing.T) { } } -// TestCreateCustomModule verifies the Create functionality -func TestCreateCustomModule(t *testing.T) { +// TestCreateSHACustomModule verifies the Create functionality +func TestCreateSHACustomModule(t *testing.T) { var buf bytes.Buffer parent := fmt.Sprintf("organizations/%s/locations/global", orgID) @@ -248,14 +253,17 @@ func TestCreateCustomModule(t *testing.T) { } } -// TestListDescendantCustomModule verifies the List Descendant functionality -func TestListDescendantCustomModule(t *testing.T) { +// TestListDescendantSHACustomModule verifies the List Descendant functionality +func TestListDescendantSHACustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + _, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestListDescendantSHACustomModule: %v", err) return } @@ -276,14 +284,17 @@ func TestListDescendantCustomModule(t *testing.T) { } } -// TestGetCustomModule verifies the Get functionality -func TestGetCustomModule(t *testing.T) { +// TestGetSHACustomModule verifies the Get functionality +func TestGetSHACustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestGetSHACustomModule: %v", err) return } @@ -305,8 +316,8 @@ func TestGetCustomModule(t *testing.T) { } } -// TestSimulateCustomModule verifies the Create functionality -func TestSimulateCustomModule(t *testing.T) { +// TestSimulateSHACustomModule verifies the Create functionality +func TestSimulateSHACustomModule(t *testing.T) { var buf bytes.Buffer parent := fmt.Sprintf("organizations/%s/locations/global", orgID) @@ -326,14 +337,17 @@ func TestSimulateCustomModule(t *testing.T) { } } -// TestListEffectiveCustomModule verifies the List Effective functionality -func TestListEffectiveCustomModule(t *testing.T) { +// TestListEffectiveSHACustomModule verifies the List Effective functionality +func TestListEffectiveSHACustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + _, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestListEffectiveSHACustomModule: %v", err) return } @@ -354,14 +368,17 @@ func TestListEffectiveCustomModule(t *testing.T) { } } -// TestUpdateCustomModule verifies the Update functionality -func TestUpdateCustomModule(t *testing.T) { +// TestUpdateSHACustomModule verifies the Update functionality +func TestUpdateSHACustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestUpdateSHACustomModule: %v", err) return } @@ -381,14 +398,17 @@ func TestUpdateCustomModule(t *testing.T) { } } -// TestGetEffectiveCustomModule verifies the Get Effective functionality -func TestGetEffectiveCustomModule(t *testing.T) { +// TestGetEffectiveSHACustomModule verifies the Get Effective functionality +func TestGetEffectiveSHACustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + createdCustomModuleID, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestGetEffectiveSHACustomModule: %v", err) return } @@ -410,14 +430,17 @@ func TestGetEffectiveCustomModule(t *testing.T) { } } -// TestListCustomModule verifies the List functionality -func TestListCustomModule(t *testing.T) { +// TestListSHACustomModule verifies the List functionality +func TestListSHACustomModule(t *testing.T) { var buf bytes.Buffer + mu.Lock() + defer mu.Unlock() + _, err := addCustomModule() if err != nil { - t.Fatalf("Could not setup test environment: %v", err) + t.Fatalf("Could not setup test environment at TestListSHACustomModule: %v", err) return }