tests: Remove global builtins state

Built ins must now be provided by the caller. We have been having a number of issues where the shared list of builtins was required to be in a different state in different tests. Related StyraInc#1101 StyraInc#1102 StyraInc#1112 StyraInc#1121 StyraInc#1129 Signed-off-by: Charlie Egan <[email protected]>
charlieegan3 · Sep 24, 2024 · 61b1193 · 61b1193
1 parent f6f2f6d
commit 61b1193
Show file tree

Hide file tree

Showing 19 changed files with 166 additions and 125 deletions.
diff --git a/internal/ast/rule.go b/internal/ast/rule.go
@@ -5,14 +5,12 @@ import (
 	"strings"
 
 	"github.com/open-policy-agent/opa/ast"
-
-	"github.com/styrainc/regal/internal/lsp/rego"
 )
 
 // GetRuleDetail returns a short descriptive string value for a given rule stating
 // if the rule is constant, multi-value, single-value etc and the type of the rule's
 // value if known.
-func GetRuleDetail(rule *ast.Rule) string {
+func GetRuleDetail(rule *ast.Rule, builtins map[string]*ast.Builtin) string {
 	if rule.Head.Args != nil {
 		return "function" + rule.Head.Args.String()
 	}
@@ -53,9 +51,7 @@ func GetRuleDetail(rule *ast.Rule) string {
 	case ast.Call:
 		name := v[0].String()
 
-		bis := rego.GetBuiltins()
-
-		if builtin, ok := bis[name]; ok {
+		if builtin, ok := builtins[name]; ok {
 			retType := builtin.Decl.NamedResult().String()
 
 			detail += fmt.Sprintf(" (%s)", simplifyType(retType))

diff --git a/internal/ast/rule_test.go b/internal/ast/rule_test.go
@@ -3,6 +3,9 @@ package ast
 import (
 	"testing"
 
+	"github.com/open-policy-agent/opa/ast"
+
+	"github.com/styrainc/regal/internal/lsp/rego"
 	"github.com/styrainc/regal/internal/parse"
 )
 
@@ -47,7 +50,9 @@ func TestGetRuleDetail(t *testing.T) {
 
 			rule := mod.Rules[0]
 
-			result := GetRuleDetail(rule)
+			bis := rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion())
+
+			result := GetRuleDetail(rule, bis)
 			if result != tc.expected {
 				t.Errorf("Expected %s, got %s", tc.expected, result)
 			}

diff --git a/internal/lsp/completions/providers/builtins.go b/internal/lsp/completions/providers/builtins.go
@@ -6,7 +6,6 @@ import (
 
 	"github.com/styrainc/regal/internal/lsp/cache"
 	"github.com/styrainc/regal/internal/lsp/hover"
-	"github.com/styrainc/regal/internal/lsp/rego"
 	"github.com/styrainc/regal/internal/lsp/types"
 	"github.com/styrainc/regal/internal/lsp/types/completion"
 )
@@ -21,7 +20,7 @@ func (*BuiltIns) Run(
 	_ context.Context,
 	c *cache.Cache,
 	params types.CompletionParams,
-	_ *Options,
+	opts *Options,
 ) ([]types.CompletionItem, error) {
 	fileURI := params.TextDocument.URI
 
@@ -45,9 +44,7 @@ func (*BuiltIns) Run(
 
 	items := []types.CompletionItem{}
 
-	bis := rego.GetBuiltins()
-
-	for _, builtIn := range bis {
+	for _, builtIn := range opts.Builtins {
 		key := builtIn.Name
 
 		if builtIn.Infix != "" {

diff --git a/internal/lsp/completions/providers/builtins_test.go b/internal/lsp/completions/providers/builtins_test.go
@@ -6,7 +6,10 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/open-policy-agent/opa/ast"
+
 	"github.com/styrainc/regal/internal/lsp/cache"
+	"github.com/styrainc/regal/internal/lsp/rego"
 	"github.com/styrainc/regal/internal/lsp/types"
 )
 
@@ -33,7 +36,9 @@ allow if c`
 		},
 	}
 
-	completions, err := p.Run(context.Background(), c, completionParams, nil)
+	completions, err := p.Run(context.Background(), c, completionParams, &Options{
+		Builtins: rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion()),
+	})
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
@@ -68,7 +73,9 @@ allow := c`
 		},
 	}
 
-	completions, err := p.Run(context.Background(), c, completionParams, nil)
+	completions, err := p.Run(context.Background(), c, completionParams, &Options{
+		Builtins: rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion()),
+	})
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
@@ -105,7 +112,9 @@ allow if {
 		},
 	}
 
-	completions, err := p.Run(context.Background(), c, completionParams, nil)
+	completions, err := p.Run(context.Background(), c, completionParams, &Options{
+		Builtins: rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion()),
+	})
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
@@ -140,7 +149,9 @@ allow if gt`
 		},
 	}
 
-	completions, err := p.Run(context.Background(), c, completionParams, nil)
+	completions, err := p.Run(context.Background(), c, completionParams, &Options{
+		Builtins: rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion()),
+	})
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
@@ -173,7 +184,9 @@ allow if c`
 		},
 	}
 
-	completions, err := p.Run(context.Background(), c, completionParams, nil)
+	completions, err := p.Run(context.Background(), c, completionParams, &Options{
+		Builtins: rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion()),
+	})
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
@@ -208,7 +221,9 @@ default allow := f`
 		},
 	}
 
-	completions, err := p.Run(context.Background(), c, completionParams, nil)
+	completions, err := p.Run(context.Background(), c, completionParams, &Options{
+		Builtins: rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion()),
+	})
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}

diff --git a/internal/lsp/completions/providers/options.go b/internal/lsp/completions/providers/options.go
@@ -1,10 +1,15 @@
 package providers
 
 import (
+	"github.com/open-policy-agent/opa/ast"
+
 	"github.com/styrainc/regal/internal/lsp/clients"
 )
 
 type Options struct {
 	ClientIdentifier clients.Identifier
 	RootURI          string
+	// Builtins is a map of built-in functions to their definitions required in
+	// the context of the current completion request.
+	Builtins map[string]*ast.Builtin
 }
diff --git a/internal/lsp/completions/providers/packagerefs_test.go b/internal/lsp/completions/providers/packagerefs_test.go
@@ -7,8 +7,11 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/open-policy-agent/opa/ast"
+
 	"github.com/styrainc/regal/internal/lsp/cache"
 	"github.com/styrainc/regal/internal/lsp/completions/refs"
+	"github.com/styrainc/regal/internal/lsp/rego"
 	"github.com/styrainc/regal/internal/lsp/types"
 	"github.com/styrainc/regal/internal/parse"
 )
@@ -42,11 +45,13 @@ import
 
 	c.SetFileContents("file:///bar/file2.rego", fileContents)
 
+	builtins := rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion())
+
 	for uri, contents := range regoFiles {
 		mod := parse.MustParseModule(contents)
 		c.SetModule(uri, mod)
 
-		c.SetFileRefs(uri, refs.DefinedInModule(mod))
+		c.SetFileRefs(uri, refs.DefinedInModule(mod, builtins))
 	}
 
 	p := &PackageRefs{}
@@ -116,11 +121,13 @@ import
 
 	c.SetFileContents("file:///file.rego", fileContents)
 
+	builtins := rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion())
+
 	for uri, contents := range regoFiles {
 		mod := parse.MustParseModule(contents)
 		c.SetModule(uri, mod)
 
-		c.SetFileRefs(uri, refs.DefinedInModule(mod))
+		c.SetFileRefs(uri, refs.DefinedInModule(mod, builtins))
 	}
 
 	p := &PackageRefs{}

diff --git a/internal/lsp/completions/providers/rulehead_test.go b/internal/lsp/completions/providers/rulehead_test.go
@@ -5,8 +5,11 @@ import (
 	"slices"
 	"testing"
 
+	"github.com/open-policy-agent/opa/ast"
+
 	"github.com/styrainc/regal/internal/lsp/cache"
 	"github.com/styrainc/regal/internal/lsp/completions/refs"
+	"github.com/styrainc/regal/internal/lsp/rego"
 	"github.com/styrainc/regal/internal/lsp/types"
 	"github.com/styrainc/regal/internal/parse"
 )
@@ -37,6 +40,8 @@ funckyfunc := true
 `,
 	}
 
+	builtins := rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion())
+
 	for uri, contents := range regoFiles {
 		mod, err := parse.Module(uri, contents)
 		if err != nil {
@@ -45,7 +50,7 @@ funckyfunc := true
 
 		c.SetFileContents(uri, contents)
 		c.SetModule(uri, mod)
-		c.SetFileRefs(uri, refs.DefinedInModule(mod))
+		c.SetFileRefs(uri, refs.DefinedInModule(mod, builtins))
 	}
 
 	p := &RuleHead{}

diff --git a/internal/lsp/completions/refs/defined.go b/internal/lsp/completions/refs/defined.go
@@ -15,7 +15,7 @@ import (
 
 // DefinedInModule returns a map of refs and details about them to be used in completions that
 // were found in the given module.
-func DefinedInModule(module *ast.Module) map[string]types.Ref {
+func DefinedInModule(module *ast.Module, builtins map[string]*ast.Builtin) map[string]types.Ref {
 	modKey := module.Package.Path.String()
 
 	// first, create a reference for the package using the metadata
@@ -92,7 +92,7 @@ func DefinedInModule(module *ast.Module) map[string]types.Ref {
 		items[ruleKey] = types.Ref{
 			Kind:        kind,
 			Label:       ruleKey,
-			Detail:      rast.GetRuleDetail(rs[0]),
+			Detail:      rast.GetRuleDetail(rs[0], builtins),
 			Description: ruleDescription,
 		}
 	}

diff --git a/internal/lsp/completions/refs/defined_test.go b/internal/lsp/completions/refs/defined_test.go
@@ -5,6 +5,9 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/open-policy-agent/opa/ast"
+
+	"github.com/styrainc/regal/internal/lsp/rego"
 	"github.com/styrainc/regal/internal/lsp/types"
 	rparse "github.com/styrainc/regal/internal/parse"
 )
@@ -32,7 +35,9 @@ func TestForModule_Package(t *testing.T) {
 package example
 `)
 
-	items := DefinedInModule(mod)
+	bis := rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion())
+
+	items := DefinedInModule(mod, bis)
 
 	expectedRefs := map[string]types.Ref{
 		"data.example": {
@@ -121,8 +126,9 @@ deny contains "strings" if true
 
 pi := 3.14
 `)
+	bis := rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion())
 
-	items := DefinedInModule(mod)
+	items := DefinedInModule(mod, bis)
 
 	expectedRefs := map[string]types.Ref{
 		"data.example": {

diff --git a/internal/lsp/documentsymbol.go b/internal/lsp/documentsymbol.go
@@ -15,6 +15,7 @@ import (
 func documentSymbols(
 	contents string,
 	module *ast.Module,
+	builtins map[string]*ast.Builtin,
 ) []types.DocumentSymbol {
 	// Only pkgSymbols would likely suffice, but we're keeping docSymbols around in case
 	// we ever want to add more top-level symbols than the package.
@@ -62,7 +63,7 @@ func documentSymbols(
 				SelectionRange: ruleRange,
 			}
 
-			if detail := rast.GetRuleDetail(rule); detail != "" {
+			if detail := rast.GetRuleDetail(rule, builtins); detail != "" {
 				ruleSymbol.Detail = &detail
 			}
 
@@ -88,7 +89,7 @@ func documentSymbols(
 				SelectionRange: groupRange,
 			}
 
-			detail := rast.GetRuleDetail(rules[0])
+			detail := rast.GetRuleDetail(rules[0], builtins)
 			if detail != "" {
 				groupSymbol.Detail = &detail
 			}
@@ -104,7 +105,7 @@ func documentSymbols(
 					SelectionRange: childRange,
 				}
 
-				childDetail := rast.GetRuleDetail(rule)
+				childDetail := rast.GetRuleDetail(rule, builtins)
 				if childDetail != "" {
 					childRule.Detail = &childDetail
 				}

diff --git a/internal/lsp/documentsymbol_test.go b/internal/lsp/documentsymbol_test.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/styrainc/regal/internal/lsp/rego"
 	"github.com/styrainc/regal/internal/lsp/types"
 	"github.com/styrainc/regal/internal/lsp/types/symbols"
 )
@@ -69,7 +70,9 @@ func TestDocumentSymbols(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			syms := documentSymbols(tc.policy, module)
+			bis := rego.BuiltinsForCapabilities(ast.CapabilitiesForThisVersion())
+
+			syms := documentSymbols(tc.policy, module, bis)
 
 			pkg := syms[0]
 			if pkg.Name != tc.expected.Name {

diff --git a/internal/lsp/hover/hover.go b/internal/lsp/hover/hover.go
@@ -135,15 +135,15 @@ func CreateHoverContent(builtin *ast.Builtin) string {
 	return result
 }
 
-func UpdateBuiltinPositions(cache *cache.Cache, uri string) error {
+func UpdateBuiltinPositions(cache *cache.Cache, uri string, builtins map[string]*ast.Builtin) error {
 	module, ok := cache.GetModule(uri)
 	if !ok {
 		return fmt.Errorf("failed to update builtin positions: no parsed module for uri %q", uri)
 	}
 
 	builtinsOnLine := map[uint][]types2.BuiltinPosition{}
 
-	for _, call := range rego.AllBuiltinCalls(module) {
+	for _, call := range rego.AllBuiltinCalls(module, builtins) {
 		line := uint(call.Location.Row)
 
 		builtinsOnLine[line] = append(builtinsOnLine[line], types2.BuiltinPosition{

diff --git a/internal/lsp/inlayhint.go b/internal/lsp/inlayhint.go
@@ -18,10 +18,10 @@ func createInlayTooltip(named *types.NamedType) string {
 	return fmt.Sprintf("%s\n\nType: `%s`", named.Descr, named.Type.String())
 }
 
-func getInlayHints(module *ast.Module) []types2.InlayHint {
+func getInlayHints(module *ast.Module, builtins map[string]*ast.Builtin) []types2.InlayHint {
 	inlayHints := make([]types2.InlayHint, 0)
 
-	for _, call := range rego.AllBuiltinCalls(module) {
+	for _, call := range rego.AllBuiltinCalls(module, builtins) {
 		for i, arg := range call.Builtin.Decl.NamedFuncArgs().Args {
 			if len(call.Args) <= i {
 				// avoid panic if provided a builtin function where the args