query-orchestration/cmd/auth_related/user.creation.tool/userlifecycle_test.go

//go:build aws

package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
	"github.com/joho/godotenv"
)

// Test configuration
const (
	testCSVPath   = "./temp-users.csv"
	testProjectID = "ebde1e1e9623491cab6f8112e67bd61c"
	testEnvID     = "9d6801123cfd4a0ea2ef1df2f430e9d3"
)

// Test data
var testUsers = []map[string]string{
	{
		"email":      "foo1@gmail.com",
		"first_name": "john",
		"last_name":  "smith",
		"roles":      "editor,viewer",
	},
	{
		"email":      "foo2@gmail.com",
		"first_name": "jane",
		"last_name":  "smith",
		"roles":      "editor,viewer,admin",
	},
}

func TestMain(m *testing.M) {
	// Load environment variables
	if err := godotenv.Overload(); err != nil {
		fmt.Println("warning - problem loading .env file - ignoring.")
	}

	// Run tests
	code := m.Run()

	// Cleanup
	if err := os.Remove(testCSVPath); err != nil && !os.IsNotExist(err) {
		fmt.Printf("Warning: failed to cleanup test CSV file: %v\n", err)
	}

	os.Exit(code)
}

func TestUserLifecycle(t *testing.T) {
	// Skip test if required environment variables are not set
	if !checkRequiredEnvVars(t) {
		return
	}

	// Create test CSV file
	if err := createTestCSV(); err != nil {
		t.Fatalf("Failed to create test CSV: %v", err)
	}

	// Test user creation
	t.Run("CreateUsers", func(t *testing.T) {
		testCreateUsers(t)
	})

	// Test user verification
	t.Run("VerifyUsersExist", func(t *testing.T) {
		testVerifyUsersExist(t)
	})

	// Test user deletion
	t.Run("DeleteUsers", func(t *testing.T) {
		testDeleteUsers(t)
	})

	// Test user deletion verification
	t.Run("VerifyUsersDeleted", func(t *testing.T) {
		testVerifyUsersDeleted(t)
	})
}

func checkRequiredEnvVars(t *testing.T) bool {
	required := []string{"PERMIT_KEY", "COGNITO_USER_POOL_ID", "AWS_REGION"}
	missing := []string{}

	for _, env := range required {
		if os.Getenv(env) == "" {
			missing = append(missing, env)
		}
	}

	if len(missing) > 0 {
		t.Skipf("Skipping integration test - missing required environment variables: %v", missing)
		return false
	}

	return true
}

func createTestCSV() error {
	csvContent := "email,first_name,last_name,role1,role2,role3\n"
	csvContent += "foo1@gmail.com,john,smith,editor,viewer,\n"
	csvContent += "foo2@gmail.com,jane,smith,editor,viewer,admin\n"

	return os.WriteFile(testCSVPath, []byte(csvContent), 0600)
}

func testCreateUsers(t *testing.T) {
	// Save original os.Args
	originalArgs := os.Args

	// Set up args for user creation
	os.Args = []string{
		"cognito-permit-sync",
		"-project", testProjectID,
		"-env", testEnvID,
		"-csv", testCSVPath,
	}

	// Restore os.Args after test
	defer func() {
		os.Args = originalArgs
	}()

	// Reset flags for testing
	resetFlags()

	// Run the application (this will create its own audit logger internally)
	err := run()
	if err != nil {
		t.Fatalf("Failed to create users: %v", err)
	}

	t.Logf("Successfully created users")
}

func testVerifyUsersExist(t *testing.T) {
	ctx := t.Context()

	// Initialize clients
	cognitoClient, err := initializeCognitoClient()
	if err != nil {
		t.Fatalf("Failed to initialize Cognito client: %v", err)
	}

	httpClient := &http.Client{Timeout: 30 * time.Second}
	config := AppConfig{
		PermitAPIKey:      os.Getenv("PERMIT_KEY"),
		PermitProjectID:   testProjectID,
		PermitEnvID:       testEnvID,
		PermitBaseURL:     "https://api.permit.io",
		CognitoUserPoolID: os.Getenv("COGNITO_USER_POOL_ID"),
		CognitoRegion:     os.Getenv("AWS_REGION"),
	}

	for _, user := range testUsers {
		t.Run(fmt.Sprintf("VerifyUser_%s", user["email"]), func(t *testing.T) {
			// Verify user exists in Cognito
			cognitoSubjectID, err := getCognitoUserSubject(ctx, cognitoClient, config.CognitoUserPoolID, user["email"])
			if err != nil {
				t.Fatalf("Failed to find user %s in Cognito: %v", user["email"], err)
			}
			if cognitoSubjectID == "" {
				t.Fatalf("User %s not found in Cognito", user["email"])
			}
			t.Logf("User %s found in Cognito with subject ID: %s", user["email"], cognitoSubjectID)

			// Verify user exists in Permit.io
			permitUserKey, err := findPermitUserByEmail(httpClient, config, user["email"])
			if err != nil {
				t.Fatalf("Failed to find user %s in Permit.io: %v", user["email"], err)
			}
			if permitUserKey == "" {
				t.Fatalf("User %s not found in Permit.io", user["email"])
			}
			t.Logf("User %s found in Permit.io with key: %s", user["email"], permitUserKey)

			// Verify subject IDs match
			if cognitoSubjectID != permitUserKey {
				t.Fatalf("Subject ID mismatch for user %s: Cognito=%s, Permit.io=%s",
					user["email"], cognitoSubjectID, permitUserKey)
			}
			t.Logf("Subject IDs match for user %s: %s", user["email"], cognitoSubjectID)

			// Verify roles are assigned in Permit.io
			if user["roles"] != "" {
				expectedRoles := strings.Split(user["roles"], ",")
				assignedRoles, err := getPermitUserRoles(httpClient, config, permitUserKey)
				if err != nil {
					t.Fatalf("Failed to get roles for user %s: %v", user["email"], err)
				}

				for _, expectedRole := range expectedRoles {
					expectedRole = strings.TrimSpace(expectedRole)
					if expectedRole == "" {
						continue
					}
					found := false
					for _, assignedRole := range assignedRoles {
						if assignedRole == expectedRole {
							found = true
							break
						}
					}
					if !found {
						t.Fatalf("Expected role %s not found for user %s. Assigned roles: %v",
							expectedRole, user["email"], assignedRoles)
					}
				}
				t.Logf("All expected roles verified for user %s: %v", user["email"], expectedRoles)
			}
		})
	}
}

func testDeleteUsers(t *testing.T) {
	// Save original os.Args
	originalArgs := os.Args

	// Set up args for user deletion
	os.Args = []string{
		"cognito-permit-sync",
		"-project", testProjectID,
		"-env", testEnvID,
		"-csv", testCSVPath,
		"--delete-users",
	}

	// Restore os.Args after test
	defer func() {
		os.Args = originalArgs
	}()

	// Reset flags for testing
	resetFlags()

	// Run the application in delete mode
	err := run()
	if err != nil {
		t.Fatalf("Failed to delete users: %v", err)
	}

	t.Logf("Successfully deleted users")
}

func testVerifyUsersDeleted(t *testing.T) {
	ctx := t.Context()

	// Initialize clients
	cognitoClient, err := initializeCognitoClient()
	if err != nil {
		t.Fatalf("Failed to initialize Cognito client: %v", err)
	}

	httpClient := &http.Client{Timeout: 30 * time.Second}
	config := AppConfig{
		PermitAPIKey:      os.Getenv("PERMIT_KEY"),
		PermitProjectID:   testProjectID,
		PermitEnvID:       testEnvID,
		PermitBaseURL:     "https://api.permit.io",
		CognitoUserPoolID: os.Getenv("COGNITO_USER_POOL_ID"),
		CognitoRegion:     os.Getenv("AWS_REGION"),
	}

	for _, user := range testUsers {
		t.Run(fmt.Sprintf("VerifyUserDeleted_%s", user["email"]), func(t *testing.T) {
			// Verify user is deleted from both systems with retry logic
			const maxRetries = 12 // 1 minute with 5-second intervals
			const retryInterval = 5 * time.Second

			var cognitoDeleted, permitDeleted bool

			for attempt := 0; attempt < maxRetries; attempt++ {
				if attempt > 0 {
					t.Logf("Retry attempt %d/%d for user %s", attempt+1, maxRetries, user["email"])
					time.Sleep(retryInterval)
				}

				// Check Cognito deletion
				if !cognitoDeleted {
					cognitoSubjectID, err := getCognitoUserSubject(ctx, cognitoClient, config.CognitoUserPoolID, user["email"])
					cognitoDeleted = (err != nil || cognitoSubjectID == "")
				}

				// Check Permit.io deletion
				if !permitDeleted {
					permitUserKey, err := findPermitUserByEmail(httpClient, config, user["email"])
					permitDeleted = (err != nil || permitUserKey == "")
				}

				// If both are deleted, we're done
				if cognitoDeleted && permitDeleted {
					break
				}
			}

			// Final verification
			if !cognitoDeleted {
				cognitoSubjectID, _ := getCognitoUserSubject(ctx, cognitoClient, config.CognitoUserPoolID, user["email"])
				t.Fatalf("User %s still exists in Cognito after %d attempts with subject ID: %s",
					user["email"], maxRetries, cognitoSubjectID)
			}
			t.Logf("User %s successfully deleted from Cognito", user["email"])

			if !permitDeleted {
				permitUserKey, _ := findPermitUserByEmail(httpClient, config, user["email"])
				t.Fatalf("User %s still exists in Permit.io after %d attempts with key: %s",
					user["email"], maxRetries, permitUserKey)
			}
			t.Logf("User %s successfully deleted from Permit.io", user["email"])
		})
	}
}

// Helper function to get Cognito user subject ID
func getCognitoUserSubject(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) (string, error) {
	input := &cognitoidentityprovider.AdminGetUserInput{
		UserPoolId: aws.String(userPoolID),
		Username:   aws.String(username),
	}

	result, err := client.AdminGetUser(ctx, input)
	if err != nil {
		// Check if it's a UserNotFoundException
		if strings.Contains(err.Error(), "UserNotFoundException") ||
			strings.Contains(err.Error(), "User does not exist") {
			return "", nil // User not found, return empty string
		}
		return "", err
	}

	// Extract subject ID from user attributes
	for _, attr := range result.UserAttributes {
		if aws.ToString(attr.Name) == "sub" {
			return aws.ToString(attr.Value), nil
		}
	}

	return "", fmt.Errorf("subject ID not found for user %s", username)
}

// Helper function to reset flags for testing
func resetFlags() {
	// Reset the default command line flag set
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
}