Repository: datreeio/datree
Pull request #94 (2021-07-11 15:40:32)

feat: Policy flag


Patch:
diff --git a/bl/evaluation/evaluator.go b/bl/evaluation/evaluator.go
index 7766f03..8269baa 100644
--- a/bl/evaluation/evaluator.go
+++ b/bl/evaluation/evaluator.go
@@ -34,10 +34,11 @@ type EvaluationResults struct {
 	}
 }
 
-func (e *Evaluator) CreateEvaluation(cliId string, cliVersion string, k8sVersion string) (*cliClient.CreateEvaluationResponse, error) {
+func (e *Evaluator) CreateEvaluation(cliId string, cliVersion string, k8sVersion string, policyName string) (*cliClient.CreateEvaluationResponse, error) {
 	createEvaluationResponse, err := e.cliClient.CreateEvaluation(&cliClient.CreateEvaluationRequest{
 		K8sVersion: &k8sVersion,
 		CliId:      cliId,
+		PolicyName: policyName,
 		Metadata: &cliClient.Metadata{
 			CliVersion:      cliVersion,
 			Os:              e.osInfo.OS,
diff --git a/bl/evaluation/printer.go b/bl/evaluation/printer.go
index 0fb3942..a17c8e1 100644
--- a/bl/evaluation/printer.go
+++ b/bl/evaluation/printer.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/datreeio/datree/bl/validation"
 
@@ -25,14 +26,14 @@ type FormattedOutput struct {
 	InvalidK8sFiles   []*validation.InvalidK8sFile
 }
 
-func PrintResults(results *EvaluationResults, invalidYamlFiles []*validation.InvalidYamlFile, invalidK8sFiles []*validation.InvalidK8sFile, evaluationSummary printer.EvaluationSummary, loginURL string, outputFormat string, printer Printer, k8sVersion string) error {
+func PrintResults(results *EvaluationResults, invalidYamlFiles []*validation.InvalidYamlFile, invalidK8sFiles []*validation.InvalidK8sFile, evaluationSummary printer.EvaluationSummary, loginURL string, outputFormat string, printer Printer, k8sVersion string, policyName string) error {
 	switch {
 	case outputFormat == "json":
 		return jsonOutput(&FormattedOutput{EvaluationResults: results, EvaluationSummary: evaluationSummary, InvalidYamlFiles: invalidYamlFiles, InvalidK8sFiles: invalidK8sFiles})
 	case outputFormat == "yaml":
 		return yamlOutput(&FormattedOutput{EvaluationResults: results, EvaluationSummary: evaluationSummary, InvalidYamlFiles: invalidYamlFiles, InvalidK8sFiles: invalidK8sFiles})
 	default:
-		return textOutput(results, invalidYamlFiles, invalidK8sFiles, evaluationSummary, loginURL, printer, k8sVersion)
+		return textOutput(results, invalidYamlFiles, invalidK8sFiles, evaluationSummary, loginURL, printer, k8sVersion, policyName)
 	}
 }
 
@@ -58,7 +59,7 @@ func yamlOutput(formattedOutput *FormattedOutput) error {
 	return nil
 }
 
-func textOutput(results *EvaluationResults, invalidYamlFiles []*validation.InvalidYamlFile, invalidK8sFiles []*validation.InvalidK8sFile, evaluationSummary printer.EvaluationSummary, url string, printer Printer, k8sVersion string) error {
+func textOutput(results *EvaluationResults, invalidYamlFiles []*validation.InvalidYamlFile, invalidK8sFiles []*validation.InvalidK8sFile, evaluationSummary printer.EvaluationSummary, url string, printer Printer, k8sVersion string, policyName string) error {
 	pwd, err := os.Getwd()
 	if err != nil {
 		return err
@@ -72,7 +73,7 @@ func textOutput(results *EvaluationResults, invalidYamlFiles []*validation.Inval
 
 	printer.PrintWarnings(warnings)
 
-	summary := parseEvaluationResultsToSummary(results, evaluationSummary, url)
+	summary := parseEvaluationResultsToSummary(results, evaluationSummary, url, policyName)
 
 	printer.PrintEvaluationSummary(evaluationSummary, k8sVersion)
 
@@ -135,8 +136,7 @@ func parseToPrinterWarnings(results *EvaluationResults, invalidYamlFiles []*vali
 type OutputTitle int
 
 const (
-	EnabledRules OutputTitle = iota
-	EvaluatedConfigurations
+	EvaluatedConfigurations OutputTitle = iota
 	TotalRulesEvaluated
 	SeeAll
 	TotalRulesPassed
@@ -145,14 +145,20 @@ const (
 
 func (t OutputTitle) String() string {
 	return [...]string{
-		"Enabled rules in policy “default”",
 		"Configs tested against policy",
 		"Total rules evaluated",
 		"See all rules in policy",
 		"Total rules passed",
 		"Total rules failed"}[t]
 }
-func parseEvaluationResultsToSummary(results *EvaluationResults, evaluationSummary printer.EvaluationSummary, loginURL string) printer.Summary {
+
+func buildEnabledRulesTitle(policyName string) string {
+	var str strings.Builder
+	fmt.Fprintf(&str, "Enabled rules in policy “%s”", policyName)
+	return str.String()
+}
+
+func parseEvaluationResultsToSummary(results *EvaluationResults, evaluationSummary printer.EvaluationSummary, loginURL string, policyName string) printer.Summary {
 	configsCount := evaluationSummary.ConfigsCount
 	totalRulesEvaluated := 0
 	totalFailedRules := 0
@@ -165,7 +171,7 @@ func parseEvaluationResultsToSummary(results *EvaluationResults, evaluationSumma
 	}
 
 	plainRows := []printer.SummaryItem{
-		{LeftCol: EnabledRules.String(), RightCol: fmt.Sprint(evaluationSummary.RulesCount), RowIndex: 0},
+		{LeftCol: buildEnabledRulesTitle(policyName), RightCol: fmt.Sprint(evaluationSummary.RulesCount), RowIndex: 0},
 		{LeftCol: EvaluatedConfigurations.String(), RightCol: fmt.Sprint(configsCount), RowIndex: 1},
 		{LeftCol: TotalRulesEvaluated.String(), RightCol: fmt.Sprint(totalRulesEvaluated), RowIndex: 2},
 		{LeftCol: SeeAll.String(), RightCol: loginURL, RowIndex: 5},
diff --git a/pkg/cliClient/evaluation.go b/pkg/cliClient/evaluation.go
index 9121cc2..4f2944c 100644
--- a/pkg/cliClient/evaluation.go
+++ b/pkg/cliClient/evaluation.go
@@ -18,12 +18,14 @@ type CreateEvaluationRequest struct {
 	CliId      string    `json:"cliId"`
 	Metadata   *Metadata `json:"metadata"`
 	K8sVersion *string   `json:"k8sVersion"`
+	PolicyName string    `json:"policyName"`
 }
 
 type CreateEvaluationResponse struct {
 	EvaluationId int    `json:"evaluationId"`
 	K8sVersion   string `json:"k8sVersion"`
 	RulesCount   int    `json:"rulesCount"`
+	PolicyName   string `json:"policyName"`
 }
 
 func (c *CliClient) CreateEvaluation(request *CreateEvaluationRequest) (*CreateEvaluationResponse, error) {
diff --git a/pkg/httpClient/client.go b/pkg/httpClient/client.go
index d2a97ab..a0814d4 100644
--- a/pkg/httpClient/client.go
+++ b/pkg/httpClient/client.go
@@ -72,7 +72,13 @@ func (c *Client) Request(method string, resourceURI string, body interface{}, he
 	}
 
 	if response.StatusCode > 399 {
-		return responseBody, fmt.Errorf("Error: %v", string(b))
+		var errorJson map[string]interface{}
+		err = json.Unmarshal(b, &errorJson)
+		if err != nil {
+			return responseBody, err
+		}
+
+		return responseBody, fmt.Errorf(fmt.Sprintf("%v", errorJson["message"]))
 	}
 
 	responseBody = Response{
diff --git a/pkg/printer/printer.go b/pkg/printer/printer.go
index 8f1a49c..ee4ed64 100644
--- a/pkg/printer/printer.go
+++ b/pkg/printer/printer.go
@@ -164,6 +164,8 @@ func (p *Printer) printInColor(title string, color *color.Color) {
 
 func (p *Printer) createNewColor(clr string) *color.Color {
 	switch clr {
+	case "error":
+		return p.theme.Colors.Error
 	case "red":
 		return p.theme.Colors.Red
 	case "yellow":
diff --git a/pkg/printer/theme.go b/pkg/printer/theme.go
index 3a9bd94..5977e2e 100644
--- a/pkg/printer/theme.go
+++ b/pkg/printer/theme.go
@@ -13,6 +13,7 @@ type theme struct {
 		Yellow *color.Color
 		Red    *color.Color
 		White  *color.Color
+		Error  *color.Color
 	}
 	Spacing struct {
 		Default string
@@ -30,10 +31,12 @@ func createTheme() *theme {
 			Yellow *color.Color
 			Red    *color.Color
 			White  *color.Color
+			Error  *color.Color
 		}{
 			Green:  color.New(color.FgGreen),
 			Yellow: color.New(color.FgYellow),
 			Red:    color.New(color.FgHiRed, color.Bold),
+			Error:  color.New(color.FgHiRed),
 			White:  color.New(color.FgHiWhite),
 		},
 		Spacing: struct{ Default string }{


Test patch:
diff --git a/bl/evaluation/evaluator_test.go b/bl/evaluation/evaluator_test.go
index 378395f..c08d795 100644
--- a/bl/evaluation/evaluator_test.go
+++ b/bl/evaluation/evaluator_test.go
@@ -77,11 +77,12 @@ func TestCreateEvaluation(t *testing.T) {
 		cliId := "test_token"
 		cliVersion := "0.0.7"
 		k8sVersion := "1.18.1"
+		policyName := "Default"
 
 		mockedCliClient.On("CreateEvaluation", mock.Anything).Return(&cliClient.CreateEvaluationResponse{EvaluationId: 1, K8sVersion: k8sVersion}, nil)
 
 		expectedCreateEvaluationResponse := &cliClient.CreateEvaluationResponse{EvaluationId: 1, K8sVersion: k8sVersion}
-		createEvaluationResponse, _ := evaluator.CreateEvaluation(cliId, cliVersion, k8sVersion)
+		createEvaluationResponse, _ := evaluator.CreateEvaluation(cliId, cliVersion, k8sVersion, policyName)
 
 		mockedCliClient.AssertCalled(t, "CreateEvaluation", mock.Anything)
 		assert.Equal(t, expectedCreateEvaluationResponse, createEvaluationResponse)
diff --git a/bl/evaluation/printer_test.go b/bl/evaluation/printer_test.go
index d6ac9c5..ef5c4cf 100644
--- a/bl/evaluation/printer_test.go
+++ b/bl/evaluation/printer_test.go
@@ -55,7 +55,7 @@ func TestPrintResults(t *testing.T) {
 		mockedPrinter.On("PrintEvaluationSummary", mock.Anything, mock.Anything)
 
 		t.Run(tt.name, func(t *testing.T) {
-			_ = PrintResults(tt.args.results, tt.args.invalidYamlFiles, tt.args.invalidK8sFiles, tt.args.evaluationSummary, tt.args.loginURL, tt.args.outputFormat, mockedPrinter, "1.18.0")
+			_ = PrintResults(tt.args.results, tt.args.invalidYamlFiles, tt.args.invalidK8sFiles, tt.args.evaluationSummary, tt.args.loginURL, tt.args.outputFormat, mockedPrinter, "1.18.0", "Default")
 
 			if tt.args.outputFormat == "json" {
 				mockedPrinter.AssertNotCalled(t, "PrintWarnings")
diff --git a/cmd/test/main.go b/cmd/test/main.go
index 799fa65..bd33e6d 100644
--- a/cmd/test/main.go
+++ b/cmd/test/main.go
@@ -22,7 +22,7 @@ import (
 
 type Evaluator interface {
 	Evaluate(filesConfigurationsChan []*extractor.FileConfigurations, evaluationId int) (*evaluation.EvaluationResults, error)
-	CreateEvaluation(cliId string, cliVersion string, k8sVersion string) (*cliClient.CreateEvaluationResponse, error)
+	CreateEvaluation(cliId string, cliVersion string, k8sVersion string, policyName string) (*cliClient.CreateEvaluationResponse, error)
 	UpdateFailedYamlValidation(invalidFiles []*validation.InvalidYamlFile, evaluationId int, stopEvaluation bool) error
 	UpdateFailedK8sValidation(invalidFiles []*validation.InvalidK8sFile, evaluationId int, stopEvaluation bool) error
 }
@@ -40,6 +40,7 @@ type TestCommandFlags struct {
 	Output               string
 	K8sVersion           string
 	IgnoreMissingSchemas bool
+	PolicyName           string
 }
 
 type EvaluationPrinter interface {
@@ -84,7 +85,7 @@ func New(ctx *TestCommandContext) *cobra.Command {
 			var err error = nil
 			defer func() {
 				if err != nil {
-					fmt.Println(err.Error())
+					ctx.Printer.PrintMessage(strings.Join([]string{"\n", err.Error(), "\n"}, ""), "error")
 				}
 			}()
 
@@ -103,7 +104,12 @@ func New(ctx *TestCommandContext) *cobra.Command {
 				return err
 			}
 
-			testCommandFlags := TestCommandFlags{Output: outputFlag, K8sVersion: k8sVersion, IgnoreMissingSchemas: ignoreMissingSchemas}
+			policy, err := cmd.Flags().GetString("policy")
+			if err != nil {
+				return err
+			}
+
+			testCommandFlags := TestCommandFlags{Output: outputFlag, K8sVersion: k8sVersion, IgnoreMissingSchemas: ignoreMissingSchemas, PolicyName: policy}
 
 			err = test(ctx, args, testCommandFlags)
 			if err != nil {
@@ -117,6 +123,7 @@ func New(ctx *TestCommandContext) *cobra.Command {
 
 	testCommand.Flags().StringP("output", "o", "", "Define output format")
 	testCommand.Flags().StringP("schema-version", "s", "", "Set kubernetes version to validate against. Defaults to 1.18.0")
+	testCommand.Flags().StringP("policy", "p", "", "Policy name to run against")
 	testCommand.Flags().BoolP("ignore-missing-schemas", "", false, "Ignore missing schemas when executing schema validation step")
 	return testCommand
 }
@@ -176,7 +183,7 @@ func test(ctx *TestCommandContext, paths []string, flags TestCommandFlags) error
 
 	validYamlFilesConfigurationsChan, invalidYamlFilesChan := files.ExtractFilesConfigurations(filesPaths, concurrency)
 
-	createEvaluationResponse, err := ctx.Evaluator.CreateEvaluation(localConfigContent.CliId, ctx.CliVersion, flags.K8sVersion)
+	createEvaluationResponse, err := ctx.Evaluator.CreateEvaluation(localConfigContent.CliId, ctx.CliVersion, flags.K8sVersion, flags.PolicyName)
 	if err != nil {
 		return err
 	}
@@ -236,7 +243,7 @@ func test(ctx *TestCommandContext, paths []string, flags TestCommandFlags) error
 		PassedPolicyCheckCount:    passedPolicyCheckCount,
 	}
 
-	err = evaluation.PrintResults(results, invalidYamlFiles, invalidK8sFiles, evaluationSummary, fmt.Sprintf("https://app.datree.io/login?cliId=%s", localConfigContent.CliId), flags.Output, ctx.Printer, createEvaluationResponse.K8sVersion)
+	err = evaluation.PrintResults(results, invalidYamlFiles, invalidK8sFiles, evaluationSummary, fmt.Sprintf("https://app.datree.io/login?cliId=%s", localConfigContent.CliId), flags.Output, ctx.Printer, createEvaluationResponse.K8sVersion, createEvaluationResponse.PolicyName)
 
 	var invocationFailedErr error = nil
 
diff --git a/cmd/test/main_test.go b/cmd/test/main_test.go
index 571657e..858477c 100644
--- a/cmd/test/main_test.go
+++ b/cmd/test/main_test.go
@@ -1,6 +1,7 @@
 package test
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/datreeio/datree/bl/validation"
@@ -23,8 +24,11 @@ func (m *mockEvaluator) Evaluate(filesConfigurationsChan []*extractor.FileConfig
 	return args.Get(0).(*evaluation.EvaluationResults), args.Error(1)
 }
 
-func (m *mockEvaluator) CreateEvaluation(cliId string, cliVersion string, k8sVersion string) (*cliClient.CreateEvaluationResponse, error) {
-	args := m.Called(cliId, cliVersion, k8sVersion)
+func (m *mockEvaluator) CreateEvaluation(cliId string, cliVersion string, k8sVersion string, policyName string) (*cliClient.CreateEvaluationResponse, error) {
+	fmt.Println("*****************")
+	fmt.Println(cliId, cliVersion, k8sVersion, policyName)
+	fmt.Println("*****************")
+	args := m.Called(cliId, cliVersion, k8sVersion, policyName)
 	return args.Get(0).(*cliClient.CreateEvaluationResponse), args.Error(1)
 }
 
@@ -121,7 +125,7 @@ func TestTestCommand(t *testing.T) {
 
 	mockedEvaluator := &mockEvaluator{}
 	mockedEvaluator.On("Evaluate", mock.Anything, mock.Anything, mock.Anything).Return(evaluationResults, nil)
-	mockedEvaluator.On("CreateEvaluation", mock.Anything, mock.Anything, mock.Anything).Return(&cliClient.CreateEvaluationResponse{EvaluationId: evaluationId, K8sVersion: "1.18.0", RulesCount: 21}, nil)
+	mockedEvaluator.On("CreateEvaluation", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&cliClient.CreateEvaluationResponse{EvaluationId: evaluationId, K8sVersion: "1.18.0", RulesCount: 21}, nil)
 	mockedEvaluator.On("UpdateFailedYamlValidation", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	mockedEvaluator.On("UpdateFailedK8sValidation", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	messager := &mockMessager{}
@@ -165,8 +169,9 @@ func TestTestCommand(t *testing.T) {
 }
 
 func test_testCommand_no_flags(t *testing.T, evaluator *mockEvaluator, filesConfigurations []*extractor.FileConfigurations, evaluationId int, ctx *TestCommandContext) {
-	test(ctx, []string{"8/*"}, TestCommandFlags{K8sVersion: "", Output: ""})
+	test(ctx, []string{"8/*"}, TestCommandFlags{K8sVersion: "1.18.0", Output: "", PolicyName: "Default"})
 
+	evaluator.AssertCalled(t, "CreateEvaluation", "134kh", "", "1.18.0", "Default")
 	evaluator.AssertCalled(t, "Evaluate", filesConfigurations, evaluationId)
 }
 
diff --git a/pkg/cliClient/cliClient_test.go b/pkg/cliClient/cliClient_test.go
index c654b24..b138a3f 100644
--- a/pkg/cliClient/cliClient_test.go
+++ b/pkg/cliClient/cliClient_test.go
@@ -321,6 +321,7 @@ func test_createEvaluation_success() *CreateEvaluationTestCase {
 			createEvaluationRequest: &CreateEvaluationRequest{
 				K8sVersion: &k8sVersion,
 				CliId:      "cli_id",
+				PolicyName: "Default",
 				Metadata: &Metadata{
 					CliVersion:      "0.0.1",
 					Os:              "darwin",
@@ -366,6 +367,7 @@ func test_createEvaluation_success() *CreateEvaluationTestCase {
 				body: &CreateEvaluationRequest{
 					K8sVersion: &k8sVersion,
 					CliId:      "cli_id",
+					PolicyName: "Default",
 					Metadata: &Metadata{
 						CliVersion:      "0.0.1",
 						Os:              "darwin",