From 4d0353f6239f828a50617772619b1a98a80386e3 Mon Sep 17 00:00:00 2001
From: Pengyuan Zhao <PengyuanZhao@users.noreply.github.com>
Date: Tue, 5 Mar 2024 10:45:33 -0500
Subject: [PATCH] test: set LW_API_TOKEN for integration tests (#1577)

* test: set LW_API_TOKEN for integration tests

* test: fetch accee_token in NewLaceworkCLI
---
 cli/cmd/cli_state.go          |  2 ++
 integration/framework_test.go | 23 +++++++++++++++++++----
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/cli/cmd/cli_state.go b/cli/cmd/cli_state.go
index 30a3003cc..ec7499932 100644
--- a/cli/cmd/cli_state.go
+++ b/cli/cmd/cli_state.go
@@ -245,6 +245,8 @@ func (c *cliState) NewClient() error {
 			api.WithTokenAndExpiration(c.Token, c.tokenCache.ExpiresAt))
 	} else if c.Token != "" {
 		apiOpts = append(apiOpts, api.WithToken(c.Token))
+	} else if os.Getenv("LW_API_TOKEN") != "" {
+		apiOpts = append(apiOpts, api.WithToken(os.Getenv("LW_API_TOKEN")))
 	}
 
 	apiOpts = append(apiOpts,
diff --git a/integration/framework_test.go b/integration/framework_test.go
index bda7bb4a2..c7d79b722 100644
--- a/integration/framework_test.go
+++ b/integration/framework_test.go
@@ -54,6 +54,7 @@ var (
 	tfPath   string
 	tf       *tfexec.Terraform
 	execPath string
+	token    string
 )
 
 // Use this function to execute a real lacework CLI command, under the hood the function
@@ -116,6 +117,24 @@ func NewLaceworkCLI(workingDir string, stdin io.Reader, args ...string) *exec.Cm
 		if os.Getenv(ciTestingUpdaterEnv) == "" {
 			env = append(env, fmt.Sprintf("%s=1", lwupdater.DisableEnv))
 		}
+		// add unique environment variable to notify the CLI that
+		// it is being executed to run our integration test suite
+		env = append(env, "LW_CLI_INTEGRATION_MODE=true")
+
+		if token == "" {
+			fmt.Println("Generating access token")
+			lacework, err := laceworkIntegrationTestClient()
+			if err != nil {
+				log.Fatal(err)
+			}
+			response, err := lacework.GenerateToken()
+			if err != nil {
+				log.Fatal(err)
+			}
+			token = response.Token
+			env = append(env, fmt.Sprintf("LW_API_TOKEN=%s", response.Token))
+		}
+
 		cmd.Env = env
 	}
 	return cmd
@@ -150,10 +169,6 @@ func runLaceworkCLI(workingDir string, args ...string) (stdout bytes.Buffer, std
 	// what they expect
 	cmd.Env = append(cmd.Env, "LW_NONINTERACTIVE=false")
 
-	// add unique environment variable to notify the CLI that
-	// it is being executed to run our integration test suite
-	cmd.Env = append(cmd.Env, "LW_CLI_INTEGRATION_MODE=true")
-
 	exitcode, err := runLaceworkCLIFromCmd(cmd)
 	if exitcode == 999 {
 		fmt.Println(stderr)