test: add TestGraphQLParseSchemaAlloc() to better test larger GraphQL schemas

Signed-off-by: Rob Myers <[email protected]>

Author: robmyersrobmyers

Diff for: v1/topdown/graphql_test.go

@@ -8,6 +8,8 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"runtime"
+	"strings"
 	"testing"
 
 	"github.com/open-policy-agent/opa/v1/ast"
@@ -889,3 +891,76 @@ func TestGraphQLParseSchema(t *testing.T) {
 		})
 	}
 }
+
+func TestGraphQLParseSchemaAlloc(t *testing.T) {
+	cases := []struct {
+		note     string
+		schema   *ast.Term
+		maxAlloc uint64
+	}{
+		{
+			note:     "default schema",
+			schema:   ast.NewTerm(ast.String(schemaWithExtraEmployeeTypes(0))),
+			maxAlloc: 1 * 1024 * 1024,
+		},
+		// Uncomment when https://github.com/open-policy-agent/opa/pull/7509 is merged
+		// {
+		// 	note:     "default schema plus 100 additional types",
+		// 	schema:   ast.NewTerm(ast.String(schemaWithExtraEmployeeTypes(100))),
+		// 	maxAlloc: 10 * 1024 * 1024,
+		// },
+		// {
+		// 	note:     "default schema plus 1,000 additional types",
+		// 	schema:   ast.NewTerm(ast.String(schemaWithExtraEmployeeTypes(1000))),
+		// 	maxAlloc: 100 * 1024 * 1024,
+		// },
+		// {
+		// 	note:     "default schema plus 10,000 additional types",
+		// 	schema:   ast.NewTerm(ast.String(schemaWithExtraEmployeeTypes(10000))),
+		// 	maxAlloc: 1000000,
+		// },
+	}
+
+	for _, tc := range cases {
+
+		t.Run(tc.note, func(t *testing.T) {
+
+			var startMemStats runtime.MemStats
+			runtime.ReadMemStats(&startMemStats)
+
+			_ = builtinGraphQLParseSchema(
+				BuiltinContext{
+					InterQueryBuiltinValueCache: nil,
+				},
+				[]*ast.Term{tc.schema},
+				func(term *ast.Term) error {
+					return nil
+				},
+			)
+
+			var finishMemStats runtime.MemStats
+			runtime.ReadMemStats(&finishMemStats)
+			allocDifference := finishMemStats.Alloc - startMemStats.Alloc
+			runtime.GC()
+
+			if allocDifference > tc.maxAlloc {
+				t.Errorf("Parsing schema '%s' expected alloc < %d, got %d", tc.note, tc.maxAlloc, allocDifference)
+				return
+			}
+		})
+	}
+}
+
+// Inflate GraphQL schema size with `count` extra types
+func schemaWithExtraEmployeeTypes(count int) string {
+
+	// build up `count` more types on basic schema
+	var builder strings.Builder
+	builder.WriteString(employeeGQLSchema)
+
+	for range count {
+		fmt.Fprintf(&builder, "\ntype Employee%d {\n    id: String!\n    salary: Int!\n}\n", count)
+	}
+
+	return builder.String()
+}