refactor: use options pattern to instantiate tests (open-policy-agent#2949)

Author: Serializator

cmd/gator/test/test.go

@@ -88,19 +88,34 @@ func run(_ *cobra.Command, _ []string) {
 		cmdutils.ErrFatalf("no input data identified")
 	}
 
-	responses, err := test.Test(unstrucs, test.Opts{
-		IncludeTrace: flagIncludeTrace,
-		GatherStats: flagGatherStats,
-		UseK8sCEL: flagEnableK8sCel,
-		Verbose: flagVerbose,
-	})
+	var printBuf bytes.Buffer
+
+	var opts []test.Opt
+	if flagIncludeTrace {
+		opts = append(opts, test.WithTrace())
+	}
+	if flagGatherStats {
+		opts = append(opts, test.WithGatherStats())
+	}
+	if flagEnableK8sCel {
+		opts = append(opts, test.WithK8sCEL(flagGatherStats))
+	}
+	if flagVerbose {
+		opts = append(opts, test.WithPrintHook(&printBuf))
+	}
+
+	responses, err := test.Test(unstrucs, opts...)
 	if err != nil {
 		cmdutils.ErrFatalf("auditing objects: %v", err)
 	}
 	results := responses.Results()
 
 	fmt.Print(formatOutput(flagOutput, results, responses.StatsEntries))
 
+	if printBuf.Len() > 0 {
+		fmt.Printf("\n%s\n", printBuf.String())
+	}
+
 	// Whether or not we return non-zero depends on whether we have a `deny`
 	// enforcementAction on one of the violated constraints
 	exitCode := 0

pkg/gator/test/test.go

@@ -3,7 +3,7 @@ package test
 import (
 	"context"
 	"fmt"
-	"os"
+	"io"
 
 	"github.com/open-policy-agent/frameworks/constraint/pkg/apis"
 	constraintclient "github.com/open-policy-agent/frameworks/constraint/pkg/client"
@@ -40,20 +40,28 @@ type Opts struct {
 	Verbose      bool
 }
 
-func Test(objs []*unstructured.Unstructured, tOpts Opts) (*GatorResponses, error) {
+type Opt func() ([]constraintclient.Opt, []rego.Arg, error)
+
+func Test(objs []*unstructured.Unstructured, opts ...Opt) (*GatorResponses, error) {
 	args := []constraintclient.Opt{constraintclient.Targets(&target.K8sValidationTarget{})}
-	if tOpts.UseK8sCEL {
-		k8sDriver, err := makeK8sCELDriver(tOpts)
+
+	driverArgs := []rego.Arg{}
+
+	for _, opt := range opts {
+		extraArgs, extraDriverArgs, err := opt()
 		if err != nil {
-			return nil, fmt.Errorf("creating K8s native driver: %w", err)
+			return nil, err
 		}
-		args = append(args, constraintclient.Driver(k8sDriver))
+
+		args = append(args, extraArgs...)
+		driverArgs = append(driverArgs, extraDriverArgs...)
 	}
 
-	driver, err := makeRegoDriver(tOpts)
+	driver, err := rego.New(driverArgs...)
 	if err != nil {
 		return nil, fmt.Errorf("creating Rego driver: %w", err)
 	}
+
 	args = append(args, constraintclient.Driver(driver), constraintclient.EnforcementPoints(util.GatorEnforcementPoint))
 
 	client, err := constraintclient.NewClient(args...)
@@ -179,30 +187,46 @@ func Test(objs []*unstructured.Unstructured, tOpts Opts) (*GatorResponses, error
 	return responses, nil
 }
 
-func makeRegoDriver(tOpts Opts) (*rego.Driver, error) {
-	var args []rego.Arg
-	if tOpts.GatherStats {
-		args = append(args, rego.GatherStats())
-	}
-	if tOpts.IncludeTrace {
-		args = append(args, rego.Tracing(tOpts.IncludeTrace))
+func WithGatherStats() Opt {
+	return func() ([]constraintclient.Opt, []rego.Arg, error) {
+		return []constraintclient.Opt{}, []rego.Arg{
+			rego.GatherStats(),
+		}, nil
 	}
-	if tOpts.Verbose {
-		args = append(args,
-			rego.PrintEnabled(tOpts.Verbose),
-			rego.PrintHook(gator.NewPrintHook(os.Stdout)),
-		)
+}
+
+func WithTrace() Opt {
+	return func() ([]constraintclient.Opt, []rego.Arg, error) {
+		return []constraintclient.Opt{}, []rego.Arg{
+			rego.Tracing(true),
+		}, nil
 	}
+}
 
-	return rego.New(args...)
+func WithPrintHook(w io.Writer) Opt {
+	return func() ([]constraintclient.Opt, []rego.Arg, error) {
+		return []constraintclient.Opt{}, []rego.Arg{
+			rego.PrintEnabled(true),
+			rego.PrintHook(gator.NewPrintHook(w)),
+		}, nil
+	}
 }
 
-func makeK8sCELDriver(tOpts Opts) (*k8scel.Driver, error) {
-	var args []k8scel.Arg
+func WithK8sCEL(gatherStats bool) Opt {
+	return func() ([]constraintclient.Opt, []rego.Arg, error) {
+		var args []k8scel.Arg
 
-	if tOpts.GatherStats {
-		args = append(args, k8scel.GatherStats())
-	}
+		if gatherStats {
+			args = append(args, k8scel.GatherStats())
+		}
 
-	return k8scel.New(args...)
+		k8sDriver, err := k8scel.New(args...)
+		if err != nil {
+			return []constraintclient.Opt{}, []rego.Arg{}, fmt.Errorf("creating K8s native driver: %w", err)
+		}
+
+		return []constraintclient.Opt{
+			constraintclient.Driver(k8sDriver),
+		}, []rego.Arg{}, nil
+	}
 }