From ef808de2738e75f7d4b4bc9876f69922b67e3809 Mon Sep 17 00:00:00 2001
From: Matt Cadorette <matthew.cadorette@lacework.net>
Date: Mon, 22 Apr 2024 09:14:17 -0400
Subject: [PATCH] fix: policy severity test correct test criteria (#1614)

---
 integration/policy_test.go | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/integration/policy_test.go b/integration/policy_test.go
index f6b6bf2e8..22384f834 100644
--- a/integration/policy_test.go
+++ b/integration/policy_test.go
@@ -21,6 +21,7 @@ package integration
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
@@ -326,12 +327,31 @@ func TestPolicyBadSeverity(t *testing.T) {
 
 }
 
+type PolicyDetail struct {
+	Severity string `json:"severity"`
+	PolicyID string `json:"policyId"`
+}
+
 func TestPolicySeverityCritical(t *testing.T) {
-	out, err, exitcode := LaceworkCLIWithTOMLConfig("policy", "list", "--severity", "critical")
-	assert.Contains(t, out.String(), "lacework-global-8")
-	assert.NotContains(t, out.String(), "high")
+	out, err, exitcode := LaceworkCLIWithTOMLConfig("policy", "list", "--severity", "critical", "--json")
 	assert.Empty(t, err.String(), "STDERR should be empty")
 	assert.Equal(t, 0, exitcode, "EXITCODE is not the expected one")
+
+	var details []PolicyDetail
+	assert.NoError(t, json.Unmarshal(out.Bytes(), &details))
+
+	found := false
+	notCritical := false
+	for _, policy := range details {
+		if policy.PolicyID == "lacework-global-8" {
+			found = true
+		}
+		if policy.Severity != "critical" {
+			notCritical = true
+		}
+	}
+	assert.True(t, found, "lacework-global-8 should have been found in policy list of severity critical")
+	assert.False(t, notCritical, "only policies with severity critical should have been found")
 }
 
 func TestPolicyShowHelp(t *testing.T) {