add support for handling unknown expression ids in test result (#1183)

aakash070 · web-flow · commit c55aebf5346f · 2025-06-03T09:16:28.000-07:00
diff --git a/tools/celtest/test_runner.go b/tools/celtest/test_runner.go
@@ -21,6 +21,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"slices"
 	"strings"
 	"testing"
 
@@ -133,7 +134,7 @@ func TriggerTests(t *testing.T, testRunnerOpts ...TestRunnerOption) {
 	if err != nil {
 		t.Fatalf("error creating test runner: %v", err)
 	}
-	programs, err := tr.Programs(t)
+	programs, err := tr.Programs(t, tr.testProgramOptions...)
 	if err != nil {
 		t.Fatalf("error creating programs: %v", err)
 	}
@@ -275,6 +276,7 @@ func DefaultTestSuiteParser(path string) TestRunnerOption {
 // - Test Suite File Path: The path to the test suite file.
 // - File Descriptor Set Path: The path to the file descriptor set file.
 // - test Suite Parser: A parser for a test suite file serialized in Textproto/YAML format.
+// - test Program Options: A list of options to be used when creating the CEL programs.
 //
 // The TestRunner provides the following methods:
 // - Programs: Creates a list of CEL programs from the input expressions.
@@ -286,6 +288,7 @@ type TestRunner struct {
 	TestSuiteFilePath     string
 	FileDescriptorSetPath string
 	testSuiteParser       TestSuiteParser
+	testProgramOptions    []cel.ProgramOption
 }
 
 // Test represents a single test case to be executed. It encompasses the following:
@@ -295,12 +298,12 @@ type TestRunner struct {
 // returns a TestResult.
 type Test struct {
 	name          string
-	input         interpreter.Activation
+	input         interpreter.PartialActivation
 	resultMatcher func(ref.Val, error) TestResult
 }
 
 // NewTest creates a new Test with the provided name, input and result matcher.
-func NewTest(name string, input interpreter.Activation, resultMatcher func(ref.Val, error) TestResult) *Test {
+func NewTest(name string, input interpreter.PartialActivation, resultMatcher func(ref.Val, error) TestResult) *Test {
 	return &Test{
 		name:          name,
 		input:         input,
@@ -417,6 +420,17 @@ func fileDescriptorSet(path string) (*descpb.FileDescriptorSet, error) {
 	return fds, nil
 }
 
+// PartialEvalProgramOption returns a TestRunnerOption which enables partial evaluation for the CEL
+// program by setting the OptPartialEval eval option.
+//
+// Note: The test setup uses env.PartialVars() for creating PartialActivation.
+func PartialEvalProgramOption() TestRunnerOption {
+	return func(tr *TestRunner) (*TestRunner, error) {
+		tr.testProgramOptions = append(tr.testProgramOptions, cel.EvalOptions(cel.OptPartialEval))
+		return tr, nil
+	}
+}
+
 // Program represents the result of creating CEL programs for the configured expressions in the
 // test runner. It encompasses the following:
 // - CELProgram - the evaluable CEL program
@@ -461,6 +475,8 @@ func (tr *TestRunner) Programs(t *testing.T, opts ...cel.ProgramOption) ([]Progr
 
 // Tests creates a list of tests from the test suite file and test suite parser configured in the
 // test runner.
+//
+// Note: The test setup uses env.PartialVars() for creating PartialActivation.
 func (tr *TestRunner) Tests(t *testing.T) ([]*Test, error) {
 	if tr.Compiler == nil {
 		return nil, fmt.Errorf("compiler is not set")
@@ -507,13 +523,14 @@ func (tr *TestRunner) createTestsFromTextproto(t *testing.T, testSuite *conforma
 	return tests, nil
 }
 
-func (tr *TestRunner) createTestInputFromPB(t *testing.T, testCase *conformancepb.TestCase) (interpreter.Activation, error) {
+func (tr *TestRunner) createTestInputFromPB(t *testing.T, testCase *conformancepb.TestCase) (interpreter.PartialActivation, error) {
 	t.Helper()
 	input := map[string]any{}
 	e, err := tr.CreateEnv()
 	if err != nil {
 		return nil, err
 	}
+	var activation interpreter.Activation
 	if testCase.GetInputContext() != nil {
 		if len(testCase.GetInput()) != 0 {
 			return nil, fmt.Errorf("only one of input and input_context can be provided at a time")
@@ -529,15 +546,22 @@ func (tr *TestRunner) createTestInputFromPB(t *testing.T, testCase *conformancep
 			if err != nil {
 				return nil, fmt.Errorf("context variable is not a valid proto: %w", err)
 			}
-			return cel.ContextProtoVars(ctx.(proto.Message))
+			activation, err = cel.ContextProtoVars(ctx.(proto.Message))
+			if err != nil {
+				return nil, fmt.Errorf("cel.ContextProtoVars() failed: %w", err)
+			}
 		case *conformancepb.InputContext_ContextMessage:
 			refVal := e.CELTypeAdapter().NativeToValue(testInput.ContextMessage)
 			ctx, err := refVal.ConvertToNative(reflect.TypeOf((*proto.Message)(nil)).Elem())
 			if err != nil {
 				return nil, fmt.Errorf("context variable is not a valid proto: %w", err)
 			}
-			return cel.ContextProtoVars(ctx.(proto.Message))
+			activation, err = cel.ContextProtoVars(ctx.(proto.Message))
+			if err != nil {
+				return nil, fmt.Errorf("cel.ContextProtoVars() failed: %w", err)
+			}
 		}
+		return e.PartialVars(activation)
 	}
 	for k, v := range testCase.GetInput() {
 		switch v.GetKind().(type) {
@@ -553,7 +577,11 @@ func (tr *TestRunner) createTestInputFromPB(t *testing.T, testCase *conformancep
 			}
 		}
 	}
-	return interpreter.NewActivation(input)
+	activation, err = interpreter.NewActivation(input)
+	if err != nil {
+		return nil, fmt.Errorf("interpreter.NewActivation(%q) failed: %w", input, err)
+	}
+	return e.PartialVars(activation)
 }
 
 func (tr *TestRunner) createResultMatcherFromPB(t *testing.T, testCase *conformancepb.TestCase) (func(ref.Val, error) TestResult, error) {
@@ -627,11 +655,34 @@ func (tr *TestRunner) createResultMatcherFromPB(t *testing.T, testCase *conforma
 			return failureResult
 		}, nil
 	case *conformancepb.TestOutput_Unknown:
-		//  TODO: to implement
+		// Validate that all expected unknown expression ids are returned by the evaluation result.
+		return func(out ref.Val, err error) TestResult {
+			expectedUnknownIDs := testOutput.Unknown.GetExprs()
+			if err == nil && types.IsUnknown(out) {
+				actualUnknownIDs := out.Value().(*types.Unknown).IDs()
+				return compareUnknownIDs(expectedUnknownIDs, actualUnknownIDs)
+			}
+			return TestResult{Success: false, Wanted: fmt.Sprintf("unknown value %v", expectedUnknownIDs), Error: err}
+		}, nil
 	}
 	return nil, nil
 }
 
+func compareUnknownIDs(expectedUnknownIDs, actualUnknownIDs []int64) TestResult {
+	sortOption := cmp.Transformer("Sort", func(in []int64) []int64 {
+		out := append([]int64{}, in...)
+		slices.Sort(out)
+		return out
+	})
+	if diff := cmp.Diff(expectedUnknownIDs, actualUnknownIDs, sortOption); diff != "" {
+		return TestResult{
+			Success: false,
+			Wanted:  fmt.Sprintf("unknown value %v", expectedUnknownIDs),
+			Error:   fmt.Errorf("mismatched test output with diff (-got +want):\n%s", diff)}
+	}
+	return TestResult{Success: true}
+}
+
 func refValueToExprValue(refVal ref.Val) (*exprpb.ExprValue, error) {
 	if types.IsUnknown(refVal) {
 		return &exprpb.ExprValue{
@@ -704,8 +755,13 @@ func (tr *TestRunner) createTestsFromYAML(t *testing.T, testSuite *test.Suite) (
 	return tests, nil
 }
 
-func (tr *TestRunner) createTestInput(t *testing.T, testCase *test.Case) (interpreter.Activation, error) {
+func (tr *TestRunner) createTestInput(t *testing.T, testCase *test.Case) (interpreter.PartialActivation, error) {
 	t.Helper()
+	e, err := tr.CreateEnv()
+	if err != nil {
+		return nil, err
+	}
+	var activation interpreter.Activation
 	if testCase.InputContext != nil && testCase.InputContext.ContextExpr != "" {
 		if len(testCase.Input) != 0 {
 			return nil, fmt.Errorf("only one of input and input_context can be provided at a time")
@@ -719,7 +775,11 @@ func (tr *TestRunner) createTestInput(t *testing.T, testCase *test.Case) (interp
 		if err != nil {
 			return nil, fmt.Errorf("context variable is not a valid proto: %w", err)
 		}
-		return cel.ContextProtoVars(ctx.(proto.Message))
+		activation, err = cel.ContextProtoVars(ctx.(proto.Message))
+		if err != nil {
+			return nil, fmt.Errorf("cel.ContextProtoVars() failed: %w", err)
+		}
+		return e.PartialVars(activation)
 	}
 	input := map[string]any{}
 	for k, v := range testCase.Input {
@@ -733,7 +793,11 @@ func (tr *TestRunner) createTestInput(t *testing.T, testCase *test.Case) (interp
 		}
 		input[k] = v.Value
 	}
-	return interpreter.NewActivation(input)
+	activation, err = interpreter.NewActivation(input)
+	if err != nil {
+		return nil, fmt.Errorf("interpreter.NewActivation(%q) failed: %w", input, err)
+	}
+	return e.PartialVars(activation)
 }
 
 func (tr *TestRunner) createResultMatcher(t *testing.T, testOutput *test.Output) (func(ref.Val, error) TestResult, error) {
@@ -793,7 +857,13 @@ func (tr *TestRunner) createResultMatcher(t *testing.T, testOutput *test.Output)
 		}, nil
 	}
 	if testOutput.UnknownSet != nil {
-		//  TODO: to implement
+		return func(out ref.Val, err error) TestResult {
+			if err == nil && types.IsUnknown(out) {
+				unknownIDs := out.Value().(*types.Unknown).IDs()
+				return compareUnknownIDs(testOutput.UnknownSet, unknownIDs)
+			}
+			return TestResult{Success: false, Wanted: fmt.Sprintf("unknown value %v", testOutput.UnknownSet), Error: err}
+		}, nil
 	}
 	return nil, nil
 }
diff --git a/tools/celtest/test_runner_test.go b/tools/celtest/test_runner_test.go
@@ -67,7 +67,7 @@ func setupTests() []*testCase {
 		},
 		{
 			name:          "raw expression test",
-			celExpression: "i + fn(j) == 42",
+			celExpression: "a || i + fn(j) == 42",
 			testSuitePath: "testdata/raw_expr_tests.yaml",
 			configPath:    "testdata/config.yaml",
 			opts:          []any{fnEnvOption()},
@@ -204,6 +204,7 @@ func TestTriggerTests(t *testing.T) {
 				DefaultTestSuiteParser(tc.testSuitePath),
 				AddFileDescriptorSet(tc.fileDescriptorSetPath),
 				TestExpression(tc.celExpression),
+				PartialEvalProgramOption(),
 			)
 			TriggerTests(t, testOpts...)
 		})
diff --git a/tools/celtest/testdata/config.yaml b/tools/celtest/testdata/config.yaml
@@ -18,3 +18,5 @@ variables:
     type_name: "int"
   - name: "j"
     type_name: "int"
+  - name: "a"
+    type_name: "bool"
diff --git a/tools/celtest/testdata/raw_expr.cel b/tools/celtest/testdata/raw_expr.cel
@@ -1 +1 @@
-i + fn(j) == 42
+a || i + fn(j) == 42
diff --git a/tools/celtest/testdata/raw_expr_tests.yaml b/tools/celtest/testdata/raw_expr_tests.yaml
@@ -22,6 +22,8 @@ section:
             value: 21
           j:
             value: 42
+          a:
+            value: false
         output:
           value: true
       - name: "false"
@@ -30,5 +32,40 @@ section:
             value: 22
           j:
             value: 42
+          a:
+            value: false
         output:
           value: false
+      - name: "true a"
+        input:
+          a:
+            value: true
+        output:
+          value: true
+  - name: "unknown"
+    tests:
+      - name: "unknown i"
+        input:
+          j:
+            value: 42
+          a:
+            value: false
+        output:
+          unknown_set:
+            - 2
+      - name: "unknown a and j"
+        input:
+          i:
+            value: 21
+        output:
+          unknown_set:
+            - 1
+            - 5
+      - name: "unknown a and i"
+        input:
+          j:
+            value: 42
+        output:
+          unknown_set:
+            - 1
+            - 2

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-i + fn(j) == 42`
	`1`	`+a \|\| i + fn(j) == 42`