cmd+tester: Parameterized tests (#7366)

Adding the ability to parameterize Rego `test_*` rules, effectively declaring multiple "test cases" within the rule, which is then individually reported (grouped under their parent rule) when running `opa test`. Fixes: #2176 Signed-off-by: Johan Fylling <johan.dev@fylling.se>
2026-08-12 19:32:48 -06:00 · 2025-02-17 22:08:27 +01:00
parent ac4f061823
commit a6fcabbba6
18 changed files with 2173 additions and 341 deletions
@@ -93,7 +93,8 @@
    "internal": [
      "internal.member_2",
      "internal.member_3",
-      "internal.print"
+      "internal.print",
+      "internal.test_case"
    ],
    "net": [
      "net.cidr_contains",
@@ -8177,6 +8178,19 @@
    "result": {},
    "wasm": false
  },
+  "internal.test_case": {
+    "args": [
+      {
+        "type": "array[any]"
+      }
+    ],
+    "available": [
+      "edge"
+    ],
+    "introduced": "edge",
+    "result": {},
+    "wasm": false
+  },
  "intersection": {
    "args": [
      {
@@ -1617,6 +1617,20 @@
        "type": "function"
      }
    },
+    {
+      "name": "internal.test_case",
+      "decl": {
+        "args": [
+          {
+            "dynamic": {
+              "type": "any"
+            },
+            "type": "array"
+          }
+        ],
+        "type": "function"
+      }
+    },
    {
      "name": "intersection",
      "decl": {
@@ -299,6 +299,154 @@ opa test --format=json pass_fail_error_test.rego
 ]
 ```

+## Parameterized Tests and Data-driven Testing
+
+A test rule can define multiple test cases for evaluation. 
+Test cases are declared by adding their name(s) to the rule as variables in its head's reference, and are evaluated through regular enumeration.
+
+**example_test.rego**:
+
+```live:example_test_cases:module:read_only
+package example_test
+
+test_concat[note] if {
+	some note, tc in {
+		"empty + empty": {
+			"a": [],
+			"b": [],
+			"exp": [],
+		},
+		"empty + filled": {
+			"a": [],
+			"b": [1, 2],
+			"exp": [1, 2],
+		},
+		"filled + filled": {
+			"a": [1, 2],
+			"b": [3, 4],
+			"exp": [1, 2, 3], # Faulty expectation, this test case will fail
+		},
+	}
+
+	act := array.concat(tc.a, tc.b)
+	act == tc.exp
+}
+```
+
+```console
+$ opa test example_test.rego
+example_test.rego:
+data.example_test.test_concat: FAIL (263.375µs)
+  empty + empty: PASS
+  empty + filled: PASS
+  filled + filled: FAIL
+--------------------------------------------------------------------------------
+FAIL: 1/1
+```
+
+Just as in regular evaluation, test-case data doesn't need to be declared as inline Rego, but can be loaded from json and yaml data files:
+
+**file_example_test.rego**:
+
+```live:example_file_test_cases:module:read_only
+package example_test
+
+import data.test_cases
+
+test_concat[note] if {
+	some note, tc in test_cases
+
+	act := array.concat(tc.a, tc.b)
+	act == tc.exp
+}
+```
+
+**file_example_test.yaml**:
+
+```yaml
+test_cases:
+   empty + empty:
+      a: []
+      b: []
+      exp: []
+   empty + filled:
+      a: []
+      b: [1, 2]
+      exp: [1, 2]
+   filled + filled:
+      a: [1, 2]
+      b: [3, 4]
+      exp: [1, 2, 3] # Faulty expectation, this test case will fail
+```
+
+```console
+$ opa test file_example_test.rego file_example_test.yaml
+file_example_test.rego:
+data.example_test.test_concat: FAIL (280µs)
+  empty + empty: PASS
+  empty + filled: PASS
+  filled + filled: FAIL
+--------------------------------------------------------------------------------
+FAIL: 1/1
+```
+
+Test cases can be nested by declaring multiple test case name variables in the head reference.
+This is useful when e.g. the same set of test cases can be used for asserting the same behaviour across slightly different circumstances:
+
+**nested_example_test.rego**:
+
+```live:example_nested_test_cases:module:read_only
+package example_test
+
+test_sign_token[note][alg] if {
+	some note, tc in {
+		"claims": {
+			"claims": {"foo": "bar"},
+		},
+		"no claims": {
+			"claims": {},
+		},
+	}
+
+	some alg in [
+		"HS256",
+		"HS333", # unknown signing algorithm, this test case will fail
+		"HS512",
+	]
+
+	secret := "foobar"
+	key := base64.encode(secret)
+
+	token := io.jwt.encode_sign({
+		"typ": "JWT",
+		"alg": alg
+	}, tc.claims, {
+		"kty": "oct",
+		"k": key
+	})
+
+	[valid, _, payload] := io.jwt.decode_verify(token, {"secret": secret})
+	valid
+	payload = tc.claims
+}
+```
+
+```console
+$ opa test nested_example_test.rego
+nested_example_test.rego:
+data.example_test.test_sign_token: FAIL (1.214541ms)
+  claims: FAIL
+    HS256: PASS
+    HS333: FAIL
+    HS512: PASS
+  no claims: FAIL
+    HS256: PASS
+    HS333: FAIL
+    HS512: PASS
+--------------------------------------------------------------------------------
+FAIL: 1/1
+```
+
 ## Data and Function Mocking

 OPA's `with` keyword can be used to replace the data document or called functions with mocks.
@@ -1,8 +1,6 @@
 module github.com/open-policy-agent/opa

-go 1.22.7
-
-toolchain go1.23.1
+go 1.23.1

 require (
 	github.com/OneOfOne/xxhash v1.2.8
@@ -299,6 +299,9 @@ var DefaultBuiltins = [...]*Builtin{
 	// Printing
 	Print,
 	InternalPrint,
+
+	// Testing
+	InternalTestCase,
 }

 // BuiltinMap provides a convenient mapping of built-in names to
@@ -3163,6 +3166,11 @@ var InternalPrint = &Builtin{
 	Decl: types.NewFunction([]types.Type{types.NewArray(nil, types.NewSet(types.A))}, nil),
 }

+var InternalTestCase = &Builtin{
+	Name: "internal.test_case",
+	Decl: types.NewFunction([]types.Type{types.NewArray(nil, types.A)}, nil),
+}
+
 /**
 * Deprecated built-ins.
 */
@@ -31,7 +31,7 @@ func TestBuiltinDeclRoundtrip(t *testing.T) {

 func TestAllBuiltinsHaveDescribedArguments(t *testing.T) {
 	for _, b := range Builtins {
-		if b.deprecated || b.Infix != "" || b.Name == "print" || b.Name == "internal.print" {
+		if b.deprecated || b.Infix != "" || b.Name == "print" || b.Name == "internal.print" || b.Name == "internal.test_case" {
 			continue
 		}

@@ -4234,6 +4234,9 @@ func (f *equalityFactory) Generate(other *Term) *Expr {
 	return expr
 }

+// TODO: Move to internal package?
+const LocalVarPrefix = "__local"
+
 type localVarGenerator struct {
 	exclude VarSet
 	suffix  string
@@ -4258,7 +4261,7 @@ func newLocalVarGenerator(suffix string, node interface{}) *localVarGenerator {

 func (l *localVarGenerator) Generate() Var {
 	for {
-		result := Var("__local" + l.suffix + strconv.Itoa(l.next) + "__")
+		result := Var(LocalVarPrefix + l.suffix + strconv.Itoa(l.next) + "__")
 		l.next++
 		if !l.exclude.Contains(result) {
 			return result
@@ -64,15 +64,16 @@ type baseDocEqIndex struct {
 }

 var (
-	equalityRef      = Equality.Ref()
-	equalRef         = Equal.Ref()
-	globMatchRef     = GlobMatch.Ref()
-	internalPrintRef = InternalPrint.Ref()
+	equalityRef         = Equality.Ref()
+	equalRef            = Equal.Ref()
+	globMatchRef        = GlobMatch.Ref()
+	internalPrintRef    = InternalPrint.Ref()
+	internalTestCaseRef = InternalTestCase.Ref()
 )

 func newBaseDocEqIndex(isVirtual func(Ref) bool) *baseDocEqIndex {
 	return &baseDocEqIndex{
-		skipIndexing:   NewSet(NewTerm(internalPrintRef)),
+		skipIndexing:   NewSet(NewTerm(internalPrintRef), NewTerm(internalTestCaseRef)),
 		isVirtual:      isVirtual,
 		root:           newTrieNodeImpl(),
 		onlyGroundRefs: true,
@@ -726,7 +726,9 @@ func (p *Parser) parseRules() []*Rule {

 	// p[x] if ...  becomes a single-value rule p[x]
 	if hasIf && !usesContains && len(rule.Head.Ref()) == 2 {
-		if !rule.Head.Ref()[1].IsGround() && len(rule.Head.Args) == 0 {
+		v := rule.Head.Ref()[1]
+		_, isRef := v.Value.(Ref)
+		if (!v.IsGround() || isRef) && len(rule.Head.Args) == 0 {
 			rule.Head.Key = rule.Head.Ref()[1]
 		}

@@ -912,6 +912,8 @@ func (head *Head) DocKind() DocKind {
 			return PartialObjectDoc
 		}
 		return PartialSetDoc
+	} else if head.HasDynamicRef() {
+		return PartialObjectDoc
 	}
 	return CompleteDoc
 }
@@ -1091,8 +1093,7 @@ func (head *Head) SetLoc(loc *Location) {

 func (head *Head) HasDynamicRef() bool {
 	pos := head.Reference.Dynamic()
-	// Ref is dynamic if it has one non-constant term that isn't the first or last term or if it's a partial set rule.
-	return pos > 0 && (pos < len(head.Reference)-1 || head.RuleKind() == MultiValue)
+	return pos > 0 && (pos < len(head.Reference))
 }

 // Copy returns a deep copy of a.
@@ -4,6 +4,13 @@ const fixtureReporterVerboseBenchmark = `FAILURES
 --------------------------------------------------------------------------------
 data.foo.bar.test_corge: FAIL (0s)

+  query:1       | Fail true = false  
+
+data.foo.bar.test_cases_fail: FAIL (0s)
+
+  query:1       | Fail true = false  
+
+  two: FAIL

 SUMMARY
 --------------------------------------------------------------------------------
@@ -11,10 +18,14 @@ data.foo.bar.test_baz	    1000	       123.0 ns/op
 data.foo.bar.test_qux: ERROR (0s)
  some err
 data.foo.bar.test_corge: FAIL (0s)
+data.foo.bar.test_cases_fail: FAIL (0s)
+  one: PASS
+  two: FAIL
+data.foo.bar.test_cases_ok	    2000	        61.50 ns/op
 --------------------------------------------------------------------------------
-PASS: 1/3
-FAIL: 1/3
-ERROR: 1/3
+PASS: 2/5
+FAIL: 2/5
+ERROR: 1/5
 `

 const fixtureReporterVerboseBenchmarkShowAllocations = `FAILURES
@@ -34,6 +34,10 @@ type PrettyReporter struct {
 	BenchMarkGoBenchFormat   bool
 }

+func (r PrettyReporter) println(a ...any) {
+	_, _ = fmt.Fprintln(r.Output, a...)
+}
+
 // Report prints the test report to the reporter's output.
 func (r PrettyReporter) Report(ch chan *Result) error {

@@ -57,44 +61,50 @@ func (r PrettyReporter) Report(ch chan *Result) error {
 	}

 	if fail > 0 && (r.Verbose || r.FailureLine) {
-		fmt.Fprintln(r.Output, "FAILURES")
+		r.println("FAILURES")
 		r.hl()

 		for _, failure := range failures {
-			fmt.Fprintln(r.Output, failure)
-			if r.Verbose {
-				fmt.Fprintln(r.Output)
-				topdown.PrettyTraceWithOpts(newIndentingWriter(r.Output), failure.Trace, topdown.PrettyTraceOptions{
-					Locations:     true,
-					ExprVariables: r.LocalVars,
-				})
-			}
+			_, _ = fmt.Fprint(r.Output, failure.string(false))
+			r.println()

-			if r.FailureLine {
-				fmt.Fprintln(r.Output)
-				for i := len(failure.Trace) - 1; i >= 0; i-- {
-					e := failure.Trace[i]
-					if e.Op == topdown.FailOp && e.Location != nil && e.QueryID != 0 {
-						if expr, isExpr := e.Node.(*ast.Expr); isExpr {
-							if _, isEvery := expr.Terms.(*ast.Every); isEvery {
-								// We're interested in the failing expression inside the every body.
-								continue
+			if len(failure.SubResults) > 0 {
+				// Print trace collectively for all sub-results.
+				if err := printFailure(r.Output, failure.Trace, r.Verbose, false, r.LocalVars); err != nil {
+					return err
+				}
+
+				if r.Verbose || r.FailureLine {
+					r.println()
+				}
+
+				for fullName, sr := range failure.SubResults.Iter {
+					w := newIndentingWriter(r.Output)
+
+					if sr.Fail {
+						if len(sr.SubResults) == 0 {
+							// Print full test-case lineage for every failed leaf sub-result for readability.
+							for _, n := range fullName {
+								_, _ = fmt.Fprintf(w, "%s: %s\n", n, sr.outcome())
+								w = newIndentingWriter(w)
+							}
+
+							if err := printFailure(w, sr.Trace, false, r.FailureLine, r.LocalVars); err != nil {
+								return err
 							}
 						}
-						_, _ = fmt.Fprintf(newIndentingWriter(r.Output), "%s:%d:\n", e.Location.File, e.Location.Row)
-						if err := topdown.PrettyEvent(newIndentingWriter(r.Output, 4), e, topdown.PrettyEventOpts{PrettyVars: r.LocalVars}); err != nil {
-							return err
-						}
-						_, _ = fmt.Fprintln(r.Output)
-						break
 					}
 				}
+			} else {
+				if err := printFailure(r.Output, failure.Trace, r.Verbose, r.FailureLine, r.LocalVars); err != nil {
+					return err
+				}
 			}

-			fmt.Fprintln(r.Output)
+			r.println()
 		}

-		fmt.Fprintln(r.Output, "SUMMARY")
+		r.println("SUMMARY")
 		r.hl()
 	}

@@ -104,25 +114,32 @@ func (r PrettyReporter) Report(ch chan *Result) error {

 		if tr.Pass() && r.BenchmarkResults {
 			dirty = true
-			fmt.Fprintln(r.Output, r.fmtBenchmark(tr))
+			r.println(r.fmtBenchmark(tr))
 		} else if r.Verbose || !tr.Pass() {
 			if tr.Location != nil && tr.Location.File != lastFile {
 				if lastFile != "" {
-					fmt.Fprintln(r.Output, "")
+					r.println("")
 				}
-				fmt.Fprintf(r.Output, "%s:\n", tr.Location.File)
+				_, _ = fmt.Fprintf(r.Output, "%s:\n", tr.Location.File)
 				lastFile = tr.Location.File
 			}
+
 			dirty = true
-			fmt.Fprintln(r.Output, tr)
+			r.println(tr.string(false))
+
+			w := newIndentingWriter(r.Output)
+			if sr := tr.SubResults; len(sr) > 0 {
+				_, _ = fmt.Fprint(w, sr.string("  "))
+			}
+
 			if len(tr.Output) > 0 {
-				fmt.Fprintln(r.Output)
-				fmt.Fprintln(newIndentingWriter(r.Output), strings.TrimSpace(string(tr.Output)))
-				fmt.Fprintln(r.Output)
+				r.println()
+				_, _ = fmt.Fprintln(newIndentingWriter(r.Output), strings.TrimSpace(string(tr.Output)))
+				r.println()
 			}
 		}
 		if tr.Error != nil {
-			fmt.Fprintf(r.Output, "  %v\n", tr.Error)
+			_, _ = fmt.Fprintf(r.Output, "  %v\n", tr.Error)
 		}
 	}

@@ -134,19 +151,52 @@ func (r PrettyReporter) Report(ch chan *Result) error {
 	total := pass + fail + skip + errs

 	if pass != 0 {
-		fmt.Fprintln(r.Output, "PASS:", fmt.Sprintf("%d/%d", pass, total))
+		r.println("PASS:", fmt.Sprintf("%d/%d", pass, total))
 	}

 	if fail != 0 {
-		fmt.Fprintln(r.Output, "FAIL:", fmt.Sprintf("%d/%d", fail, total))
+		r.println("FAIL:", fmt.Sprintf("%d/%d", fail, total))
 	}

 	if skip != 0 {
-		fmt.Fprintln(r.Output, "SKIPPED:", fmt.Sprintf("%d/%d", skip, total))
+		r.println("SKIPPED:", fmt.Sprintf("%d/%d", skip, total))
 	}

 	if errs != 0 {
-		fmt.Fprintln(r.Output, "ERROR:", fmt.Sprintf("%d/%d", errs, total))
+		r.println("ERROR:", fmt.Sprintf("%d/%d", errs, total))
+	}
+
+	return nil
+}
+
+func printFailure(w io.Writer, trace []*topdown.Event, verbose bool, failureLine bool, localVars bool) error {
+	if verbose {
+		_, _ = fmt.Fprintln(w)
+		topdown.PrettyTraceWithOpts(newIndentingWriter(w), trace, topdown.PrettyTraceOptions{
+			Locations:     true,
+			ExprVariables: localVars,
+		})
+	}
+
+	if failureLine {
+		_, _ = fmt.Fprintln(w)
+		for i := len(trace) - 1; i >= 0; i-- {
+			e := trace[i]
+			if e.Op == topdown.FailOp && e.Location != nil && e.QueryID != 0 {
+				if expr, isExpr := e.Node.(*ast.Expr); isExpr {
+					if _, isEvery := expr.Terms.(*ast.Every); isEvery {
+						// We're interested in the failing expression inside the every body.
+						continue
+					}
+				}
+				_, _ = fmt.Fprintf(newIndentingWriter(w), "%s:%d:\n", e.Location.File, e.Location.Row)
+				if err := topdown.PrettyEvent(newIndentingWriter(w, 4), e, topdown.PrettyEventOpts{PrettyVars: localVars}); err != nil {
+					return err
+				}
+				_, _ = fmt.Fprintln(w)
+				break
+			}
+		}
 	}

 	return nil
@@ -253,6 +303,12 @@ func newIndentingWriter(w io.Writer, indent ...int) indentingWriter {
 	if len(indent) > 0 {
 		i = indent[0]
 	}
+
+	if iw, ok := w.(indentingWriter); ok {
+		i += iw.indent
+		w = iw.w
+	}
+
 	return indentingWriter{
 		w:      w,
 		indent: i,
@@ -92,6 +92,69 @@ func TestPrettyReporterVerbose(t *testing.T) {
 				File: "policy4.rego",
 			},
 		},
+		{
+			Package: "data.foo.qux",
+			Name:    "test_cases",
+			Trace:   getFakeTraceEvents(),
+			Fail:    true,
+			// Will be sorted to "bar", "baz", "foo" in output for stability
+			SubResults: SubResultMap{
+				"foo": {
+					Name: "foo",
+					Fail: false,
+				},
+				"bar": {
+					Name: "bar",
+					Fail: true,
+				},
+				"baz": {
+					Name: "baz",
+					Fail: false,
+				},
+			},
+			Location: &ast.Location{
+				File: "policy5.rego",
+			},
+		},
+		{
+			Package: "data.foo.qux",
+			Name:    "test_cases_nested",
+			Trace:   getFakeTraceEvents(),
+			Fail:    true,
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: false,
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: true,
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+			},
+			Location: &ast.Location{
+				File: "policy5.rego",
+			},
+		},
 	}

 	r := PrettyReporter{
@@ -110,6 +173,19 @@ data.foo.bar.test_corge: FAIL (0s)

  query:1       | Fail true = false  

+data.foo.qux.test_cases: FAIL (0s)
+
+  query:1       | Fail true = false  
+
+  bar: FAIL
+
+data.foo.qux.test_cases_nested: FAIL (0s)
+
+  query:1       | Fail true = false  
+
+  two: FAIL
+    foo: FAIL
+
 SUMMARY
 --------------------------------------------------------------------------------
 policy1.rego:
@@ -129,11 +205,24 @@ data.foo.bar.test_contains_print: PASS (0s)

 policy4.rego:
 data.foo.baz.p.q.r.test_quz: PASS (0s)
+
+policy5.rego:
+data.foo.qux.test_cases: FAIL (0s)
+  bar: FAIL
+  baz: PASS
+  foo: PASS
+data.foo.qux.test_cases_nested: FAIL (0s)
+  one: PASS
+    bar: PASS
+    foo: PASS
+  two: FAIL
+    bar: PASS
+    foo: FAIL
 --------------------------------------------------------------------------------
-PASS: 3/6
-FAIL: 1/6
-SKIPPED: 1/6
-ERROR: 1/6
+PASS: 3/8
+FAIL: 3/8
+SKIPPED: 1/8
+ERROR: 1/8
 `

 	str := buf.String()
@@ -232,6 +321,68 @@ func TestPrettyReporterFailureLine(t *testing.T) {
 				File: "policy3.rego",
 			},
 		},
+		{
+			Package: "data.foo.qux",
+			Name:    "test_cases_nested",
+			Trace:   getFakeTraceEvents(),
+			Fail:    true,
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: false,
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: true,
+							Trace: getFakeTraceEventsFor(
+								ast.MustParseExpr("x == y + z"),
+								func(e *topdown.Event) {
+									// QueryID == 0 is not pretty-printed, as this is the base query to eval the test rule; not the test rule itself.
+									e.QueryID = 1
+								},
+								func(e *topdown.Event) {
+									e.Location.File = "policy5.rego"
+									e.Location.Row = 5
+								},
+								func(e *topdown.Event) {
+									e.Locals = ast.NewValueMap()
+									e.Locals.Put(ast.Var("x"), ast.Number("1"))
+									e.Locals.Put(ast.Var("y"), ast.Number("2"))
+									e.Locals.Put(ast.Var("z"), ast.Number("3"))
+								},
+								func(e *topdown.Event) {
+									e.LocalMetadata = map[ast.Var]topdown.VarMetadata{
+										"x": {Name: "x"},
+										"y": {Name: "y"},
+										"z": {Name: "z"},
+									}
+								}),
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+			},
+			Location: &ast.Location{
+				File: "policy5.rego",
+			},
+		},
 	}

 	r := PrettyReporter{
@@ -262,6 +413,18 @@ data.foo.bar.test_contains_print_fail: FAIL (0s)
 data.foo.baz.p.q.r.test_quz: FAIL (0s)


+data.foo.qux.test_cases_nested: FAIL (0s)
+
+  two: FAIL
+    foo: FAIL
+      
+        policy5.rego:5:
+          x == y + z
+          |    |   |
+          |    |   3
+          |    2
+          1      
+
 SUMMARY
 --------------------------------------------------------------------------------
 policy1.rego:
@@ -278,11 +441,20 @@ data.foo.bar.test_contains_print_fail: FAIL (0s)

 policy3.rego:
 data.foo.baz.p.q.r.test_quz: FAIL (0s)
+
+policy5.rego:
+data.foo.qux.test_cases_nested: FAIL (0s)
+  one: PASS
+    bar: PASS
+    foo: PASS
+  two: FAIL
+    bar: PASS
+    foo: FAIL
 --------------------------------------------------------------------------------
-PASS: 2/7
-FAIL: 3/7
-SKIPPED: 1/7
-ERROR: 1/7
+PASS: 2/8
+FAIL: 4/8
+SKIPPED: 1/8
+ERROR: 1/8
 `

 	if exp != buf.String() {
@@ -357,6 +529,69 @@ func TestPrettyReporter(t *testing.T) {
 				File: "policy3.rego",
 			},
 		},
+		{
+			Package: "data.foo.qux",
+			Name:    "test_cases",
+			Trace:   getFakeTraceEvents(),
+			Fail:    true,
+			// Will be sorted to "bar", "baz", "foo" in output for stability
+			SubResults: SubResultMap{
+				"foo": {
+					Name: "foo",
+					Fail: false,
+				},
+				"bar": {
+					Name: "bar",
+					Fail: true,
+				},
+				"baz": {
+					Name: "baz",
+					Fail: false,
+				},
+			},
+			Location: &ast.Location{
+				File: "policy4.rego",
+			},
+		},
+		{
+			Package: "data.foo.qux",
+			Name:    "test_cases_nested",
+			Trace:   getFakeTraceEvents(),
+			Fail:    true,
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: false,
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: true,
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+			},
+			Location: &ast.Location{
+				File: "policy4.rego",
+			},
+		},
 	}

 	r := PrettyReporter{
@@ -382,11 +617,24 @@ data.foo.bar.test_contains_print_fail: FAIL (0s)

 policy3.rego:
 data.foo.baz.p.q.r.test_quz: FAIL (0s)
+
+policy4.rego:
+data.foo.qux.test_cases: FAIL (0s)
+  bar: FAIL
+  baz: PASS
+  foo: PASS
+data.foo.qux.test_cases_nested: FAIL (0s)
+  one: PASS
+    bar: PASS
+    foo: PASS
+  two: FAIL
+    bar: PASS
+    foo: FAIL
 --------------------------------------------------------------------------------
-PASS: 2/7
-FAIL: 3/7
-SKIPPED: 1/7
-ERROR: 1/7
+PASS: 2/9
+FAIL: 5/9
+SKIPPED: 1/9
+ERROR: 1/9
 `

 	if exp != buf.String() {
@@ -429,6 +677,45 @@ func TestJSONReporter(t *testing.T) {
 			Package: "data.foo.baz",
 			Name:    "p.q.r.test_quz",
 		},
+		{
+			Package: "data.foo.qux",
+			Name:    "test_cases_nested",
+			Trace:   getFakeTraceEvents(),
+			Fail:    true,
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: false,
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+					SubResults: SubResultMap{
+						"foo": {
+							Name: "foo",
+							Fail: true,
+						},
+						"bar": {
+							Name: "bar",
+							Fail: false,
+						},
+					},
+				},
+			},
+			Location: &ast.Location{
+				File: "policy5.rego",
+			},
+		},
 	}

 	r := JSONReporter{
@@ -441,162 +728,197 @@ func TestJSONReporter(t *testing.T) {
 		t.Fatal(err)
 	}

-	exp := util.MustUnmarshalJSON([]byte(`[
-  {
-    "location": null,
-    "package": "data.foo.bar",
-    "name": "test_baz",
-    "duration": 0,
-    "trace": [
-      {
-        "Op": "Fail",
-        "Node": {
-          "index": 0,
-          "terms": [
-            {
-              "type": "ref",
-              "value": [
-                {
-                  "type": "var",
-                  "value": "eq"
-                }
-              ]
-            },
-            {
-              "type": "boolean",
-              "value": true
-            },
-            {
-              "type": "boolean",
-              "value": false
-            }
-          ]
+	exp := util.MustUnmarshalJSON([]byte(`[ {
+  "location" : null,
+  "package" : "data.foo.bar",
+  "name" : "test_baz",
+  "duration" : 0,
+  "trace" : [ {
+    "Op" : "Fail",
+    "Node" : {
+      "index" : 0,
+      "terms" : [ {
+        "type" : "ref",
+        "value" : [ {
+          "type" : "var",
+          "value" : "eq"
+        } ]
+      }, {
+        "type" : "boolean",
+        "value" : true
+      }, {
+        "type" : "boolean",
+        "value" : false
+      } ]
+    },
+    "Location" : {
+      "file" : "",
+      "row" : 1,
+      "col" : 1
+    },
+    "QueryID" : 0,
+    "ParentID" : 0,
+    "Locals" : null,
+    "LocalMetadata" : null,
+    "Message" : "",
+    "Ref" : null
+  } ]
+}, {
+  "location" : null,
+  "package" : "data.foo.bar",
+  "name" : "test_qux",
+  "error" : { },
+  "duration" : 0,
+  "trace" : [ {
+    "Op" : "Fail",
+    "Node" : {
+      "index" : 0,
+      "terms" : [ {
+        "type" : "ref",
+        "value" : [ {
+          "type" : "var",
+          "value" : "eq"
+        } ]
+      }, {
+        "type" : "boolean",
+        "value" : true
+      }, {
+        "type" : "boolean",
+        "value" : false
+      } ]
+    },
+    "Location" : {
+      "file" : "",
+      "row" : 1,
+      "col" : 1
+    },
+    "QueryID" : 0,
+    "ParentID" : 0,
+    "Locals" : null,
+    "LocalMetadata" : null,
+    "Message" : "",
+    "Ref" : null
+  } ]
+}, {
+  "location" : null,
+  "package" : "data.foo.bar",
+  "name" : "test_corge",
+  "fail" : true,
+  "duration" : 0,
+  "trace" : [ {
+    "Op" : "Fail",
+    "Node" : {
+      "index" : 0,
+      "terms" : [ {
+        "type" : "ref",
+        "value" : [ {
+          "type" : "var",
+          "value" : "eq"
+        } ]
+      }, {
+        "type" : "boolean",
+        "value" : true
+      }, {
+        "type" : "boolean",
+        "value" : false
+      } ]
+    },
+    "Location" : {
+      "file" : "",
+      "row" : 1,
+      "col" : 1
+    },
+    "QueryID" : 0,
+    "ParentID" : 0,
+    "Locals" : null,
+    "LocalMetadata" : null,
+    "Message" : "",
+    "Ref" : null
+  } ]
+}, {
+  "location" : null,
+  "package" : "data.foo.bar",
+  "name" : "todo_test_qux",
+  "skip" : true,
+  "duration" : 0
+}, {
+  "location" : null,
+  "package" : "data.foo.bar",
+  "name" : "test_contains_print",
+  "duration" : 0,
+  "output" : "ZmFrZSBwcmludCBvdXRwdXQK"
+}, {
+  "location" : null,
+  "package" : "data.foo.baz",
+  "name" : "p.q.r.test_quz",
+  "duration" : 0
+}, {
+  "location" : {
+    "file" : "policy5.rego",
+    "row" : 0,
+    "col" : 0
+  },
+  "package" : "data.foo.qux",
+  "name" : "test_cases_nested",
+  "fail" : true,
+  "duration" : 0,
+  "trace" : [ {
+    "Op" : "Fail",
+    "Node" : {
+      "index" : 0,
+      "terms" : [ {
+        "type" : "ref",
+        "value" : [ {
+          "type" : "var",
+          "value" : "eq"
+        } ]
+      }, {
+        "type" : "boolean",
+        "value" : true
+      }, {
+        "type" : "boolean",
+        "value" : false
+      } ]
+    },
+    "Location" : {
+      "file" : "",
+      "row" : 1,
+      "col" : 1
+    },
+    "QueryID" : 0,
+    "ParentID" : 0,
+    "Locals" : null,
+    "LocalMetadata" : null,
+    "Message" : "",
+    "Ref" : null
+  } ],
+  "sub_results" : {
+    "one" : {
+      "name" : "one",
+      "sub_results" : {
+        "bar" : {
+          "name" : "bar"
        },
-        "Location": {
-          "file": "",
-          "row": 1,
-          "col": 1
-        },
-        "QueryID": 0,
-        "ParentID": 0,
-        "Locals": null,
-        "LocalMetadata": null,
-        "Message": "",
-        "Ref": null
+        "foo" : {
+          "name" : "foo"
+        }
      }
-    ]
-  },
-  {
-    "location": null,
-    "package": "data.foo.bar",
-    "name": "test_qux",
-    "error": {},
-    "duration": 0,
-    "trace": [
-      {
-        "Op": "Fail",
-        "Node": {
-          "index": 0,
-          "terms": [
-            {
-              "type": "ref",
-              "value": [
-                {
-                  "type": "var",
-                  "value": "eq"
-                }
-              ]
-            },
-            {
-              "type": "boolean",
-              "value": true
-            },
-            {
-              "type": "boolean",
-              "value": false
-            }
-          ]
+    },
+    "two" : {
+      "name" : "two",
+      "fail" : true,
+      "sub_results" : {
+        "bar" : {
+          "name" : "bar"
        },
-        "Location": {
-          "file": "",
-          "row": 1,
-          "col": 1
-        },
-        "QueryID": 0,
-        "ParentID": 0,
-        "Locals": null,
-        "LocalMetadata": null,
-        "Message": "",
-        "Ref": null
+        "foo" : {
+          "name" : "foo",
+          "fail" : true
+        }
      }
-    ]
-  },
-  {
-    "location": null,
-    "package": "data.foo.bar",
-    "name": "test_corge",
-    "fail": true,
-    "duration": 0,
-    "trace": [
-      {
-        "Op": "Fail",
-        "Node": {
-          "index": 0,
-          "terms": [
-            {
-              "type": "ref",
-              "value": [
-                {
-                  "type": "var",
-                  "value": "eq"
-                }
-              ]
-            },
-            {
-              "type": "boolean",
-              "value": true
-            },
-            {
-              "type": "boolean",
-              "value": false
-            }
-          ]
-        },
-        "Location": {
-          "file": "",
-          "row": 1,
-          "col": 1
-        },
-        "QueryID": 0,
-        "ParentID": 0,
-        "Locals": null,
-        "LocalMetadata": null,
-        "Message": "",
-        "Ref": null
-      }
-    ]
-  },
-  {
-    "location": null,
-    "package": "data.foo.bar",
-    "name": "todo_test_qux",
-    "skip": true,
-    "duration": 0
-  },
-  {
-	  "location": null,
-	  "package": "data.foo.bar",
-	  "name": "test_contains_print",
-	  "output": "ZmFrZSBwcmludCBvdXRwdXQK",
-	  "duration": 0
-  },
-  {
-	"location": null,
-	"package": "data.foo.baz",
-	"name": "p.q.r.test_quz",
-	"duration": 0
-}
-]
+    }
+  }
+} ]
 `))

 	result := util.MustUnmarshalJSON(buf.Bytes())
@@ -633,6 +955,7 @@ func TestPrettyReporterVerboseBenchmark(t *testing.T) {
 		{
 			Package: "data.foo.bar",
 			Name:    "test_corge",
+			Trace:   getFakeTraceEvents(),
 			Fail:    true,
 			BenchmarkResult: &testing.BenchmarkResult{
 				N:         100,
@@ -643,6 +966,53 @@ func TestPrettyReporterVerboseBenchmark(t *testing.T) {
 				Extra:     nil,
 			},
 		},
+		{
+			Package: "data.foo.bar",
+			Name:    "test_cases_fail",
+			Fail:    true,
+			Trace:   getFakeTraceEvents(),
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+				},
+			},
+			BenchmarkResult: &testing.BenchmarkResult{
+				N:         100,
+				T:         12300,
+				Bytes:     0,
+				MemAllocs: 567,
+				MemBytes:  890,
+				Extra:     nil,
+			},
+		},
+		{
+			Package: "data.foo.bar",
+			Name:    "test_cases_ok",
+			Fail:    false,
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+				},
+			},
+			BenchmarkResult: &testing.BenchmarkResult{
+				N:         2000,
+				T:         123000,
+				Bytes:     0,
+				MemAllocs: 567,
+				MemBytes:  890,
+				Extra:     nil,
+			},
+		},
 	}

 	r := PrettyReporter{
@@ -817,6 +1187,52 @@ func TestJSONReporterBenchmark(t *testing.T) {
 			Name:    "todo_test_qux",
 			Skip:    true,
 		},
+		{
+			Package: "data.foo.bar",
+			Name:    "test_cases_fail",
+			Fail:    true,
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+				},
+			},
+			BenchmarkResult: &testing.BenchmarkResult{
+				N:         100,
+				T:         12300,
+				Bytes:     0,
+				MemAllocs: 567,
+				MemBytes:  890,
+				Extra:     nil,
+			},
+		},
+		{
+			Package: "data.foo.bar",
+			Name:    "test_cases_ok",
+			Fail:    false,
+			SubResults: SubResultMap{
+				"one": {
+					Name: "one",
+					Fail: false,
+				},
+				"two": {
+					Name: "two",
+					Fail: true,
+				},
+			},
+			BenchmarkResult: &testing.BenchmarkResult{
+				N:         2000,
+				T:         123000,
+				Bytes:     0,
+				MemAllocs: 567,
+				MemBytes:  890,
+				Extra:     nil,
+			},
+		},
 	}

 	r := JSONReporter{
@@ -835,7 +1251,7 @@ func TestJSONReporterBenchmark(t *testing.T) {
    "package": "data.foo.bar",
    "name": "test_baz",
    "duration": 0,
-	"benchmark_result": {
+    "benchmark_result": {
      "N": 1000,
      "T": 123000,
      "Bytes": 0,
@@ -861,11 +1277,58 @@ func TestJSONReporterBenchmark(t *testing.T) {
    "duration": 0
  },
  {
-    "location":null,
-    "duration":0,
-    "name":"todo_test_qux",
-    "package":"data.foo.bar",
-    "skip":true
+    "location": null,
+    "package": "data.foo.bar",
+    "name": "todo_test_qux",
+    "skip": true,
+    "duration": 0
+  },
+  {
+    "location": null,
+    "package": "data.foo.bar",
+    "name": "test_cases_fail",
+    "fail": true,
+    "duration": 0,
+    "benchmark_result": {
+      "N": 100,
+      "T": 12300,
+      "Bytes": 0,
+      "MemAllocs": 567,
+      "MemBytes": 890,
+      "Extra": null
+    },
+    "sub_results": {
+      "one": {
+        "name": "one"
+      },
+      "two": {
+        "name": "two",
+        "fail": true
+      }
+    }
+  },
+  {
+    "location": null,
+    "package": "data.foo.bar",
+    "name": "test_cases_ok",
+    "duration": 0,
+    "benchmark_result": {
+      "N": 2000,
+      "T": 123000,
+      "Bytes": 0,
+      "MemAllocs": 567,
+      "MemBytes": 890,
+      "Extra": null
+    },
+    "sub_results": {
+      "one": {
+        "name": "one"
+      },
+      "two": {
+        "name": "two",
+        "fail": true
+      }
+    }
  }
 ]
 `))
@@ -0,0 +1,474 @@
+// Copyright 2025 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package tester
+
+import (
+	"testing"
+
+	"github.com/open-policy-agent/opa/v1/ast"
+)
+
+func TestInjectTestCaseFunc(t *testing.T) {
+	testCases := []struct {
+		note   string
+		module string
+		exp    string
+	}{
+		{
+			note: "no head-ref, assigned last in body",
+			module: `package test
+				test_foo if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					tc.x == 1
+				}`,
+			// func not injected
+			exp: `package test
+				test_foo if {
+					__local3__ = [{"note": "a", "x": 1}]
+					__local2__ = __local3__[__local1__]
+					__local2__.x = 1
+				}`,
+		},
+
+		{
+			note: "manual use of internal.test_case",
+			module: `package test
+				test_foo[tc.note] if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					tc.x == 1
+					internal.test_case([tc.note, "foo", "bar"])
+				}`,
+			// func not injected
+			exp: `package test
+				test_foo[__local0__] if { 
+					__local4__ = [{"note": "a", "x": 1}]
+					__local3__ = __local4__[__local2__]            # func would have been injected subsequent to here
+					__local3__.x = 1
+					__local5__ = __local3__.note
+					internal.test_case([__local5__, "foo", "bar"]) # manual use of func
+					__local0__ = __local3__.note 
+				}`,
+		},
+
+		{
+			note: "head-ref, assigned last in body",
+			module: `package test
+				test_foo.foo if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					tc.x == 1
+				}`,
+			// no var assignment in body, func injected first in body
+			exp: `package test
+				test_foo.foo if { 
+					internal.test_case(["foo"])          # func injection
+					__local3__ = [{"note": "a", "x": 1}]
+					__local2__ = __local3__[__local1__]
+					__local2__.x = 1
+				}`,
+		},
+		{
+			note: "string in head-ref, assigned last in body",
+			module: `package test
+				test_foo["foo"] if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					tc.x == 1
+				}`,
+			// no var assignment in body, func injected first in body
+			exp: `package test
+				test_foo.foo if { 
+					internal.test_case(["foo"])          # func injection
+					__local3__ = [{"note": "a", "x": 1}]
+					__local2__ = __local3__[__local1__]
+					__local2__.x = 1
+				}`,
+		},
+
+		{
+			note: "const in head-ref, assigned last in body",
+			module: `package test
+				foo := "bar"
+				test_foo[foo] if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					tc.x == 1
+				}`,
+			// var assignment can be moved up the body, func injected after moved expr
+			exp: `package test
+				foo := "bar" if { true }
+				test_foo[__local0__] if { 
+					__local0__ = data.test.foo           # generated head-ref const/var assignment, moved up
+					internal.test_case([__local0__])     # func injection
+					__local4__ = [{"note": "a", "x": 1}]
+					__local3__ = __local4__[__local2__]
+					__local3__.x = 1
+				}`,
+		},
+
+		{
+			note: "var in head-ref, assigned last in body",
+			module: `package test
+				test_foo[note] if {
+					some tc in [
+						{"note": "a"},
+					]
+					note := tc.note
+				}`,
+			// var assignment cannot be moved up the body, func injected last in body
+			exp: `package test
+				test_foo[__local3__] if { 
+					__local4__ = [{"note": "a"}]
+					__local2__ = __local4__[__local1__]
+					__local3__ = __local2__.note        # head-ref var assignment
+					internal.test_case([__local3__])    # func injection
+				}`,
+		},
+		{
+			note: "var in head-ref, assigned last in body, trailing assertions",
+			module: `package test
+				test_foo[note] if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					note := tc.note
+					tc.x == 1
+				}`,
+			// var assignment cannot be moved up the body, func injected last in body
+			exp: `package test
+				test_foo[__local3__] if { 
+					__local4__ = [{"note": "a", "x": 1}]
+					__local2__ = __local4__[__local1__]
+					__local3__ = __local2__.note        # head-ref var assignment
+					internal.test_case([__local3__])    # func injection
+					__local2__.x = 1
+				}`,
+		},
+		{
+			note: "var in head-ref, assigned mid-body",
+			module: `package test
+				test_foo[note] if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					note := tc.note
+					tc.x == 1
+				}`,
+			// var assignment cannot be moved up the body, func injected after assignment
+			exp: `package test
+				test_foo[__local3__] if { 
+					__local4__ = [{"note": "a", "x": 1}]
+					__local2__ = __local4__[__local1__]
+					__local3__ = __local2__.note        # head-ref var assignment
+					internal.test_case([__local3__])    # func injection
+					__local2__.x = 1
+				}`,
+		},
+		{
+			note: "var in head-ref, assigned after unrelated assertions",
+			module: `package test
+				test_foo[note] if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					tc.x == 1
+					note := tc.note
+				}`,
+			// var assignment can be moved up the body, func injected after assignment
+			exp: `package test
+				test_foo[__local3__] if { 
+					__local4__ = [{"note": "a", "x": 1}]
+					__local2__ = __local4__[__local1__]
+					__local2__.x = 1                    # non-generated expression, can't be moved
+					__local3__ = __local2__.note        # head-ref var assignment
+					internal.test_case([__local3__])    # func injection
+				}`,
+		},
+		{
+			note: "var in head-ref, non-assignment reference in body",
+			module: `package test
+				test_concat[note] if {
+					some note, tc in {     # Compiled into roughly '__local__ = {...}; tc = __local__[note]' 
+						"empty + empty": {
+							"a": [],
+							"b": [],
+							"exp": [],
+						},
+					}
+					
+					act := array.concat(tc.a, tc.b)
+					act == tc.exp
+				}`,
+			exp: `package test
+				test_concat[__local0__] if { 
+					__local3__ = {"empty + empty": {"a": [], "b": [], "exp": []}}
+					__local1__ = __local3__[__local0__] # head-ref var assignment
+					internal.test_case([__local0__])    # func injection
+					__local5__ = __local1__.a
+					__local6__ = __local1__.b
+					array.concat(__local5__, __local6__, __local4__)
+					__local2__ = __local4__
+					__local2__ = __local1__.exp
+				}`,
+		},
+
+		{
+			note: "ref in head-ref",
+			module: `package test
+				test_foo[tc.note] if {
+					some tc in [
+						{"note": "a"},
+					]
+				}`,
+			// var assignment cannot be moved up the body, func injected last in body
+			exp: `package test
+				test_foo[__local0__] if { 
+					__local4__ = [{"note": "a"}]
+					__local3__ = __local4__[__local2__]
+					__local0__ = __local3__.note        # generated head-ref var assignment
+					internal.test_case([__local0__])    # func injection
+				}`,
+		},
+		{
+			note: "ref in head-ref, can move above unrelated assertions",
+			module: `package test
+				test_foo[tc.note] if {
+					some tc in [
+						{"note": "a", "x": 1, "y": 2},
+					]
+					tc.x == 1
+					tc.y == 2
+				}`,
+			// var assignment can be moved up the body, func injected after assignment
+			exp: `package test
+				test_foo[__local0__] if { 
+					__local4__ = [{"note": "a", "x": 1, "y": 2}]
+					__local3__ = __local4__[__local2__]
+					__local0__ = __local3__.note                 # generated head-ref var assignment
+					internal.test_case([__local0__])			 # func injection
+					__local3__.x = 1
+					__local3__.y = 2
+					}`,
+		},
+
+		// Multiple head-ref test-case terms
+
+		{
+			note: "multi-term head-ref, assigned last in body",
+			module: `package test
+				test_foo.foo.bar if {
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					tc.x == 1
+				}`,
+			// no var assignment in body, func injected first in body
+			exp: `package test
+				test_foo.foo.bar if { 
+					internal.test_case(["foo", "bar"])   # func injection
+					__local3__ = [{"note": "a", "x": 1}]
+					__local2__ = __local3__[__local1__]
+					__local2__.x = 1
+				}`,
+		},
+		{
+			note: "multiple vars in head-ref",
+			module: `package test
+				test_foo[note1][note2] if {
+					some flag in [
+						{"note": "on", "a": 1},
+					]
+					note1 := flag.note
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					note2 := tc.note
+					flag.a == 1
+					tc.x == 1
+				}`,
+			// var assignment cannot be moved up the body, func injected after last head-ref assignment
+			exp: `package test
+				test_foo[__local3__][__local7__] = true if { 
+					__local8__ = [{"a": 1, "note": "on"}]; 
+					__local2__ = __local8__[__local1__]; 
+					__local3__ = __local2__.note;                 # manual head-ref var assignment
+					__local9__ = [{"note": "a", "x": 1}]; 
+					__local6__ = __local9__[__local5__]; 
+					__local7__ = __local6__.note;                 # manual head-ref var assignment
+					internal.test_case([__local3__, __local7__]); # func injection, after last head-ref assignment
+					__local2__.a = 1; 
+					__local6__.x = 1 
+				}`,
+		},
+		{
+			note: "multiple vars in head-ref, manual assignment below unrelated assertion(s)",
+			module: `package test
+				test_foo[note1][note2] if {
+					some flag in [
+						{"note": "on", "a": 1},
+					]
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					note2 := tc.note
+					flag.a == 1
+					note1 := flag.note
+					tc.x == 1
+				}`,
+			// var assignment cannot be moved up the body, func injected after last head-ref assignment
+			exp: `package test
+				test_foo[__local7__][__local6__] if { 
+					__local8__ = [{"a": 1, "note": "on"}]
+					__local2__ = __local8__[__local1__]
+					__local9__ = [{"note": "a", "x": 1}]
+					__local5__ = __local9__[__local4__]
+					__local6__ = __local5__.note                 # manual head-ref var assignment 
+					__local2__.a = 1
+					__local7__ = __local2__.note                 # manual head-ref var assignment, cannot be moved
+					internal.test_case([__local7__, __local6__]) # func injection, after last head-ref assignment
+					__local5__.x = 1 
+				}`,
+		},
+		{
+			note: "multiple refs in head-ref",
+			module: `package test
+				test_foo[tc.note][flag.note] if {
+					some flag in [
+						{"note": "on", "a": 1},
+					]
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					flag.a == 1
+					tc.x == 1
+				}`,
+			// var assignment can be moved up the body, func injected after last head-ref assignment
+			exp: `package test
+				test_foo[__local0__][__local1__] if { 
+					__local8__ = [{"a": 1, "note": "on"}]
+					__local4__ = __local8__[__local3__]
+					__local1__ = __local4__.note                 # generated head-ref var assignment, moved up
+					__local9__ = [{"note": "a", "x": 1}]
+					__local7__ = __local9__[__local6__]
+					__local0__ = __local7__.note                 # generated head-ref var assignment, moved up
+					internal.test_case([__local0__, __local1__]) # func injection, after last head-ref assignment
+					__local4__.a = 1; 
+					__local7__.x = 1 
+				}`,
+		},
+		{
+			note: "multiple refs in head-ref, mixed with ground terms",
+			module: `package test
+				test_foo[tc.note].bar[flag.note].baz if {
+					some flag in [
+						{"note": "on", "a": 1},
+					]
+					some tc in [
+						{"note": "a", "x": 1},
+					]
+					flag.a == 1
+					tc.x == 1
+				}`,
+			// var assignment cannot be moved up the body, func injected last in body
+			exp: `package test
+				test_foo[__local0__].bar[__local1__].baz if { 
+					__local8__ = [{"a": 1, "note": "on"}]
+					__local4__ = __local8__[__local3__]
+					__local1__ = __local4__.note                               # generated head-ref var assignment, moved up
+					__local9__ = [{"note": "a", "x": 1}]
+					__local7__ = __local9__[__local6__]
+					__local0__ = __local7__.note                               # generated head-ref var assignment, moved up
+					internal.test_case([__local0__, "bar", __local1__, "baz"]) # func injection, after last head-ref assignment
+					__local4__.a = 1
+					__local7__.x = 1
+				}`,
+		},
+		{
+			note: "multiple vars in head-ref, non-assignment reference in body",
+			module: `package example_test
+				test_sign_token[note][alg] if {
+					some note, tc in {
+						"claims": {
+							"claims": {"foo": "bar"},
+						},
+						"no claims": {
+							"claims": {},
+						},
+					}
+				
+					some alg in [
+						"HS256",
+						"HS512",
+					]
+				
+					secret := "foobar"
+					key := base64.encode(secret)
+				
+					token := io.jwt.encode_sign({
+						"typ": "JWT",
+						"alg": alg
+					}, tc.claims, {
+						"kty": "oct",
+						"k": key
+					})
+				
+					[valid, _, payload] := io.jwt.decode_verify(token, {"secret": secret})
+					valid
+					payload = tc.claims
+				}`,
+			exp: `package example_test
+				test_sign_token[__local0__][__local4__] if { 
+					__local11__ = {"claims": {"claims": {"foo": "bar"}}, "no claims": {"claims": {}}}
+					__local1__ = __local11__[__local0__]
+					__local12__ = ["HS256", "HS512"]
+					__local4__ = __local12__[__local3__]
+					internal.test_case([__local0__, __local4__])                  # func injection
+					__local5__ = "foobar"; base64.encode(__local5__, __local13__)
+					__local6__ = __local13__
+					__local16__ = __local1__.claims
+					io.jwt.encode_sign({"alg": __local4__, "typ": "JWT"}, __local16__, {"k": __local6__, "kty": "oct"}, __local14__)
+					__local7__ = __local14__
+					io.jwt.decode_verify(__local7__, {"secret": __local5__}, __local15__)
+					[__local8__, __local9__, __local10__] = __local15__
+					__local8__
+					__local10__ = __local1__.claims
+				}`,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.note, func(t *testing.T) {
+			modules := map[string]*ast.Module{
+				"test.rego": ast.MustParseModule(tc.module),
+			}
+
+			exp := ast.MustParseModule(tc.exp)
+
+			c := ast.NewCompiler()
+			c.WithStageAfter("RewriteLocalVars", ast.CompilerStageDefinition{
+				Name:       "InjectTestCaseFunc",
+				MetricName: "inject_test_case_func",
+				Stage:      injectTestCaseFunc,
+			})
+
+			c.Compile(modules)
+			if c.Failed() {
+				t.Fatalf("Unexpected error(s): %v", c.Errors)
+			}
+
+			result := c.Modules["test.rego"]
+			if !result.Equal(exp) {
+				t.Fatalf("Expected:\n\n%v\n\nbut got:\n\n%v", exp, result)
+			}
+		})
+	}
+}
@@ -8,9 +8,11 @@ package tester
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"regexp"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -56,6 +58,77 @@ func RunWithFilter(ctx context.Context, _ loader.Filter, paths ...string) ([]*Re
 	return result, nil
 }

+type SubResult struct {
+	Name       string           `json:"name,omitempty"`
+	Fail       bool             `json:"fail,omitempty"`
+	Trace      []*topdown.Event `json:"-"`
+	SubResults SubResultMap     `json:"sub_results,omitempty"`
+}
+
+type SubResultMap map[string]*SubResult
+
+func (srm SubResultMap) Update(path ast.Array, trace []*topdown.Event) bool {
+	strPath := make([]string, path.Len())
+	for i := range path.Len() {
+		strPath[i] = termToString(path.Elem(i))
+	}
+	return srm.update(strPath, 0, trace)
+}
+
+func (srm SubResultMap) update(path []string, i int, trace []*topdown.Event) bool {
+	if i >= len(path) {
+		return true
+	}
+
+	k := path[i]
+	entry, ok := srm[k]
+	if !ok {
+		entry = &SubResult{
+			Name:       path[i],
+			Fail:       true,
+			SubResults: SubResultMap{},
+		}
+		srm[k] = entry
+	}
+
+	if i == len(path)-1 {
+		entry.Trace = trace
+		return entry.Fail
+	}
+
+	fail := entry.SubResults.update(path, i+1, trace)
+
+	if fail {
+		entry.Fail = true
+	}
+
+	return fail
+}
+
+type unknownResolver struct{}
+
+func (unknownResolver) Resolve(_ ast.Ref) (interface{}, error) {
+	return "UNKNOWN", nil
+}
+
+func termToString(t *ast.Term) string {
+	ti, err := ast.ValueToInterface(t.Value, unknownResolver{})
+	if err != nil {
+		return "INVALID"
+	}
+	var str string
+	var ok bool
+	if str, ok = ti.(string); !ok {
+		var buf bytes.Buffer
+		if err := json.NewEncoder(&buf).Encode(ti); err != nil {
+			return "INVALID"
+		}
+		str = strings.TrimSpace(buf.String())
+	}
+
+	return str
+}
+
 // Result represents a single test case result.
 type Result struct {
 	Location        *ast.Location            `json:"location"`
@@ -69,16 +142,18 @@ type Result struct {
 	Output          []byte                   `json:"output,omitempty"`
 	FailedAt        *ast.Expr                `json:"failed_at,omitempty"`
 	BenchmarkResult *testing.BenchmarkResult `json:"benchmark_result,omitempty"`
+	SubResults      SubResultMap             `json:"sub_results,omitempty"`
 }

 func newResult(loc *ast.Location, pkg, name string, duration time.Duration, trace []*topdown.Event, output []byte) *Result {
 	return &Result{
-		Location: loc,
-		Package:  pkg,
-		Name:     name,
-		Duration: duration,
-		Trace:    trace,
-		Output:   output,
+		Location:   loc,
+		Package:    pkg,
+		Name:       name,
+		Duration:   duration,
+		Trace:      trace,
+		Output:     output,
+		SubResults: SubResultMap{},
 	}
 }

@@ -88,10 +163,23 @@ func (r Result) Pass() bool {
 }

 func (r *Result) String() string {
+	return r.string(true)
+}
+
+func (r *Result) string(subResults bool) string {
 	if r.Skip {
 		return fmt.Sprintf("%v.%v: %v", r.Package, r.Name, r.outcome())
 	}
-	return fmt.Sprintf("%v.%v: %v (%v)", r.Package, r.Name, r.outcome(), r.Duration)
+	var buf bytes.Buffer
+
+	buf.WriteString(fmt.Sprintf("%v.%v: %v (%v)", r.Package, r.Name, r.outcome(), r.Duration))
+
+	if subResults {
+		buf.WriteString("\n")
+		buf.WriteString(r.SubResults.String())
+	}
+
+	return buf.String()
 }

 func (r *Result) outcome() string {
@@ -107,6 +195,52 @@ func (r *Result) outcome() string {
 	return "ERROR"
 }

+func (sr *SubResult) String() string {
+	return fmt.Sprintf("%v: %v", sr.Name, sr.outcome())
+}
+
+func (sr *SubResult) outcome() string {
+	if sr.Fail {
+		return "FAIL"
+	}
+	return "PASS"
+}
+
+// Iter is a depth-first iterator over all sub-results.
+func (srm SubResultMap) Iter(yield func([]string, *SubResult) bool) {
+	srm.iter(nil, yield)
+}
+
+func (srm SubResultMap) iter(namePrefix []string, yield func([]string, *SubResult) bool) {
+	for _, k := range util.KeysSorted(srm) {
+		sr := srm[k]
+
+		fullName := make([]string, len(namePrefix)+1)
+		copy(fullName, namePrefix)
+		fullName[len(fullName)-1] = k
+
+		if !yield(fullName, sr) {
+			return
+		}
+		sr.SubResults.iter(fullName, yield)
+	}
+}
+
+func (srm SubResultMap) String() string {
+	return srm.string("  ")
+}
+
+func (srm SubResultMap) string(indent string) string {
+	var buf bytes.Buffer
+	for fullName, sr := range srm.Iter {
+		buf.WriteString(fmt.Sprintf("%s%s\n",
+			strings.Repeat(indent, len(fullName)-1),
+			sr.String(),
+		))
+	}
+	return buf.String()
+}
+
 // BenchmarkOptions defines options specific to benchmarking tests
 type BenchmarkOptions struct {
 	ReportAllocations bool
@@ -309,6 +443,12 @@ func (r *Runner) runTests(ctx context.Context, txn storage.Transaction, enablePr
 		Stage:      rewriteDuplicateTestNames,
 	})

+	r.compiler.WithStageAfter("RewriteLocalVars", ast.CompilerStageDefinition{
+		Name:       "InjectTestCaseFunc",
+		MetricName: "inject_test_case_func",
+		Stage:      injectTestCaseFunc,
+	})
+
 	if r.store == nil {
 		r.store = inmem.NewWithOpts(inmem.OptRoundTripOnWrite(false))
 	}
@@ -379,20 +519,33 @@ func (r *Runner) runTests(ctx context.Context, txn storage.Transaction, enablePr
 }

 func (r *Runner) shouldRun(rule *ast.Rule, testRegex *regexp.Regexp) bool {
-	ruleName := ruleName(rule.Head)
+	var ref ast.Ref

-	// All tests must have the right prefix
-	if !strings.HasPrefix(ruleName, TestPrefix) && !strings.HasPrefix(ruleName, SkipTestPrefix) {
-		return false
+	for _, term := range rule.Head.Ref().GroundPrefix() {
+		ref = ref.Append(term)
+
+		var n string
+		switch v := term.Value.(type) {
+		case ast.Var:
+			n = string(v)
+		case ast.String:
+			n = string(v)
+		default:
+			n = ""
+		}
+
+		if strings.HasPrefix(n, TestPrefix) || strings.HasPrefix(n, SkipTestPrefix) {
+			// Even with the prefix it needs to pass the regex (if applicable)
+			fullName := rule.Module.Package.Path.Extend(ref).String()
+			if testRegex != nil && !testRegex.MatchString(fullName) {
+				return false
+			}
+
+			return true
+		}
 	}

-	// Even with the prefix it needs to pass the regex (if applicable)
-	fullName := rule.Ref().String()
-	if testRegex != nil && !testRegex.MatchString(fullName) {
-		return false
-	}
-
-	return true
+	return false
 }

 // rewriteDuplicateTestNames will rewrite duplicate test names to have a numbered suffix.
@@ -402,19 +555,23 @@ func rewriteDuplicateTestNames(compiler *ast.Compiler) *ast.Error {
 	count := map[string]int{}
 	for _, mod := range compiler.Modules {
 		for _, rule := range mod.Rules {
-			name := ruleName(rule.Head)
+			name, ref := ruleName(rule.Head)
 			if !strings.HasPrefix(name, TestPrefix) {
 				continue
 			}
-			key := rule.Ref().String()
+
+			key := mod.Package.Path.Extend(ref).String()
 			if k, ok := count[key]; ok {
-				ref := rule.Head.Ref()
+				dynamicSuffix := rule.Head.Ref()[len(ref):]
 				newName := fmt.Sprintf("%s#%02d", name, k)
 				if len(ref) == 1 {
 					ref[0] = ast.VarTerm(newName)
 				} else {
 					ref[len(ref)-1] = ast.StringTerm(newName)
 				}
+				for i := range len(dynamicSuffix) {
+					ref = append(ref, dynamicSuffix[i])
+				}
 				rule.Head.SetRef(ref)
 			}
 			count[key]++
@@ -423,54 +580,296 @@ func rewriteDuplicateTestNames(compiler *ast.Compiler) *ast.Error {
 	return nil
 }

+var testCaseFuncRef = ast.InternalTestCase.Ref()
+
+// injectTestCaseFunc will inject a call to the 'internal.test_case' function into partial-object test rules.
+// We attempt to find the earliest point in the rule body where we can inject the call, to ensure that the test-case
+// function is called as early as possible so that we capture as many failed test cases as possible.
+// This may require us to move generated assignment expressions up the body.
+// We do not attempt to move non-generated expressions, as that could contradict author intent.
+//
+// Consider the test rule:
+//
+//	test_concat[tc.note] if {
+//		some tc in [{
+//			"note": "empty + empty",
+//			"a": [],
+//			"b": [],
+//			"exp": [],
+//		}]
+//		act := array.concat(tc.a, tc.b)
+//		act == tc.exp
+//	}
+//
+// The compiler will rewrite this rule to (mid-stage @ 'RewriteLocalVars'):
+//
+//	test_concat[__local0__] := true if {
+//		__local3__ = [{"a": [], "b": [], "exp": [], "note": "empty + empty"}][__local2__]
+//		__local4__ = array.concat(__local3__.a, __local3__.b)
+//		__local4__ == __local3__.exp
+//		__local0__ = __local3__.note # generated var
+//	}
+//
+// We move the generated var assignment as far up the body as possible, and inject the test-case function below it:
+//
+//	test_concat[__local0__] := true if {
+//		__local3__ = [{"a": [], "b": [], "exp": [], "note": "empty + empty"}][__local2__]
+//		__local0__ = __local3__.note                          # moved up
+//		internal.test_case([__local0__])                      # injected
+//		__local4__ = array.concat(__local3__.a, __local3__.b) # this and below expressions can now fail eval and we will still have captured the test-case
+//		__local4__ == __local3__.exp
+//	}
+func injectTestCaseFunc(compiler *ast.Compiler) *ast.Error {
+	for _, mod := range compiler.Modules {
+		for _, rule := range mod.Rules {
+			// Only apply to test rules
+			rName, rRef := ruleName(rule.Head)
+			if !strings.HasPrefix(rName, TestPrefix) {
+				continue
+			}
+
+			// Only apply to rules that doesn't have manual use of the test-case function
+			manualCall := false
+			ast.WalkExprs(rule.Body, func(expr *ast.Expr) bool {
+				if expr.IsCall() && expr.Operator().Equal(testCaseFuncRef) {
+					manualCall = true
+					return true
+				}
+				return false
+			})
+
+			if manualCall {
+				continue
+			}
+
+			// Construct test-case name
+			ref := rule.Head.Ref()
+			if len(ref) <= len(rRef) {
+				// We only inject the test-case function if there is a rule ref "tail" behind the rule name
+				continue
+			}
+			argsRef := ref[len(rRef):]
+			args := ast.NewArray(argsRef...)
+
+			//
+			// Pass 1: Move generated assignment expressions up the body
+			//
+
+			for _, term := range argsRef {
+				// We expect to find generated expressions - if any - at the tail of the body, so we start from the end
+				for i := len(rule.Body) - 1; i >= 0; {
+					expr := rule.Body[i]
+					moved := false
+
+					// If the expression is a generated assignment of a var in the head ref, we attempt to move it as far
+					// up the body as possible.
+					// This is a shallow move, we don't attempt to detect multiple levels of indirection and don't move such expressions; in such case, we move the assigning expression up to the first reference.
+					// Once done for all vars in the head ref, we can inject the test case function below the last (possibly moved) such expr.
+					// Note: We don't move non-generated expressions, as that could contradict author intent.
+					if expr.Generated && (expr.IsEquality() || expr.IsAssignment()) && expr.Operand(0).Equal(term) {
+						// Based on the vars in the rhs of the expr, see if we can move it up the rule body
+						// FIXME: Can we get away with just placing it under the lowes first occurrence of any referenced var?
+						vars := ast.NewVarSet()
+						ast.WalkVars(expr.Operand(1), func(v ast.Var) bool {
+							// We only care about local vars
+							if isLocalVar(v) {
+								vars.Add(v)
+							}
+							return false
+						})
+
+						if len(vars) == 0 {
+							// No local vars referenced, can be moved to top of body
+							rule.Body, moved = moveExpr(rule.Body, i, 0)
+						} else {
+							// Find the lowest (highest up the body) individual index of each var referenced in the rhs,
+							// and select the highest (lowest down the body) of those
+
+							// TODO: Use TypedValueMap once synced with main
+							lowest := ast.NewValueMap()
+
+							for j := i - 1; j >= 0; j-- {
+								expr := rule.Body[j]
+								ast.WalkVars(expr, func(v ast.Var) bool {
+									if vars.Contains(v) {
+										// We override the value for each var, so we get the lowest index (line highest up the body) for each
+										lowest.Put(v, ast.Number(strconv.Itoa(j)))
+										return true
+									}
+									return false
+								})
+							}
+
+							highest := 0
+							lowest.Iter(func(k, v ast.Value) bool {
+								if n, err := strconv.Atoi(string(v.(ast.Number))); err == nil {
+									if n > highest {
+										highest = n
+									}
+								}
+								return false
+							})
+
+							if highest < i {
+								// The expression is lower in the body than the lowes line of any expression that might contribute to its assignment
+								// Move the expression to just after the lowest line
+								moveTo := highest + 1
+								rule.Body, moved = moveExpr(rule.Body, i, moveTo)
+							}
+						}
+					}
+
+					// If the expression was moved, we need to re-evaluate the current index, as it contains a new expression
+					if !moved {
+						i--
+					}
+				}
+			}
+
+			//
+			// Pass 2: Inject the test-case function below the lowest first occurrence of any referenced var
+			//
+
+			injectBelowMap := ast.NewValueMap()
+			for _, term := range argsRef {
+				for i := len(rule.Body) - 1; i >= 0; i-- {
+					expr := rule.Body[i]
+
+					ast.WalkVars(expr, func(v ast.Var) bool {
+						if term.Value.Compare(v) == 0 {
+							injectBelowMap.Put(v, ast.Number(strconv.Itoa(i)))
+						}
+						return false
+					})
+				}
+			}
+
+			// Find the earliest point where the test case function can be injected
+			injectBelow := -1
+			injectBelowMap.Iter(func(k, v ast.Value) bool {
+				if n, err := strconv.Atoi(string(v.(ast.Number))); err == nil {
+					if n > injectBelow {
+						injectBelow = n
+					}
+				}
+				return false
+			})
+
+			testCaseFuncExpr := ast.NewExpr([]*ast.Term{
+				ast.NewTerm(ast.InternalTestCase.Ref()),
+				ast.NewTerm(args),
+			})
+
+			rule.Body = insertExpr(rule.Body, testCaseFuncExpr, injectBelow+1)
+		}
+	}
+	return nil
+}
+
+func isLocalVar(v ast.Value) bool {
+	if v, ok := v.(ast.Var); ok {
+		if strings.HasPrefix(string(v), ast.LocalVarPrefix) {
+			return true
+		}
+	}
+	return false
+}
+
+func insertExpr(body ast.Body, expr *ast.Expr, index int) ast.Body {
+	if index <= 0 {
+		return append(ast.Body{expr}, body...)
+	}
+
+	if index >= len(body) {
+		return append(body, expr)
+	}
+
+	return append(body[:index], append(ast.Body{expr}, body[index:]...)...)
+}
+
+func moveExpr(body ast.Body, from int, to int) (ast.Body, bool) {
+	if from == to {
+		return body, false
+	}
+
+	expr := body[from]                                                // Save the expression to move
+	body = append(body[:from], body[from+1:]...)                      // Remove the expression from the body
+	body = append(body[:to], append(ast.Body{expr}, body[to:]...)...) // Insert the expression at the new position
+	return body, true
+}
+
 // ruleName is a helper to be used when checking if a function
 // (a) is a test, or
 // (b) needs to be skipped
 // -- it'll resolve `p.q.r` to `r`. For representing results, we'll
 // use rule.Head.Ref()
-func ruleName(h *ast.Head) string {
-	ref := h.Ref()
-	switch last := ref[len(ref)-1].Value.(type) {
-	case ast.Var:
-		return string(last)
-	case ast.String:
-		return string(last)
-	default:
-		return ""
+func ruleName(h *ast.Head) (string, ast.Ref) {
+	var n string
+	var ref ast.Ref
+
+	for _, term := range h.Ref().GroundPrefix() {
+		ref = ref.Append(term)
+		switch v := term.Value.(type) {
+		case ast.Var:
+			n = string(v)
+		case ast.String:
+			n = string(v)
+		default:
+			n = ""
+		}
+
+		if strings.HasPrefix(n, TestPrefix) || strings.HasPrefix(n, SkipTestPrefix) {
+			break
+		}
 	}
+
+	return n, ref
 }

 func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.Module, rule *ast.Rule) (*Result, bool) {
-	var bufferTracer *topdown.BufferTracer
-	var tracer topdown.QueryTracer
-
-	if r.cover != nil {
-		tracer = r.cover
-	} else if r.trace {
-		bufferTracer = topdown.NewBufferTracer()
-		tracer = bufferTracer
-	}
-
-	ruleName := ruleName(rule.Head)
+	ruleName, ruleRef := ruleName(rule.Head)
 	if strings.HasPrefix(ruleName, SkipTestPrefix) { // TODO(sr): add test
-		tr := newResult(rule.Loc(), mod.Package.Path.String(), rule.Head.Ref().String(), 0*time.Second, nil, nil)
+		tr := newResult(rule.Loc(), mod.Package.Path.String(), ruleRef.String(), 0*time.Second, nil, nil)
 		tr.Skip = true
 		return tr, false
 	}

+	var bufferTracer *topdown.BufferTracer
+	var tracers []topdown.QueryTracer
+
+	if r.cover != nil {
+		t := NewTestQueryTracer()
+		tracers = append(tracers, r.cover, t)
+		bufferTracer = &t.BufferTracer
+	} else if r.trace {
+		bufferTracer = topdown.NewBufferTracer()
+		tracers = append(tracers, bufferTracer)
+	} else {
+		t := NewTestQueryTracer()
+		tracers = append(tracers, t)
+		bufferTracer = &t.BufferTracer
+	}
+
 	printbuf := bytes.NewBuffer(nil)
 	var builtinErrors []topdown.Error
-	rg := rego.New(
+	queryPath := rule.Module.Package.Path.Extend(ruleRef)
+
+	opts := []func(*rego.Rego){
 		rego.Store(r.store),
 		rego.Transaction(txn),
 		rego.Compiler(r.compiler),
-		rego.Query(rule.Path().String()),
-		rego.QueryTracer(tracer),
+		rego.Query(queryPath.String()),
 		rego.Runtime(r.runtime),
 		rego.Target(r.target),
 		rego.PrintHook(topdown.NewPrintHook(printbuf)),
 		rego.BuiltinErrorList(&builtinErrors),
-	)
+	}
+
+	for _, t := range tracers {
+		opts = append(opts, rego.QueryTracer(t))
+	}
+
+	rg := rego.New(opts...)

 	// Register custom builtins on rego instance
 	for _, v := range r.customBuiltins {
@@ -486,7 +885,7 @@ func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.
 		trace = *bufferTracer
 	}

-	tr := newResult(rule.Loc(), mod.Package.Path.String(), rule.Head.Ref().String(), dt, trace, printbuf.Bytes())
+	tr := newResult(rule.Loc(), mod.Package.Path.String(), ruleRef.String(), dt, trace, printbuf.Bytes())

 	// If there was an error other than errors from builtins, prefer that error.
 	if err != nil {
@@ -506,6 +905,8 @@ func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.
 		}
 	} else if len(rs) == 0 {
 		tr.Fail = true
+	} else if rule.Head.DocKind() == ast.PartialObjectDoc {
+		tr.Fail, tr.SubResults = subResults(rs[0].Expressions[0].Value, trace)
 	} else if b, ok := rs[0].Expressions[0].Value.(bool); !ok || !b {
 		tr.Fail = true
 	}
@@ -513,11 +914,99 @@ func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.
 	return tr, stop
 }

+func subResults(v any, trace []*topdown.Event) (bool, map[string]*SubResult) {
+	if v == nil {
+		return true, map[string]*SubResult{}
+	}
+
+	var fail bool
+	result := SubResultMap{}
+
+	switch x := v.(type) {
+	case map[string]any:
+		for k, v := range x {
+			sr := subResult(k, v)
+			result[k] = sr
+			if sr.Fail {
+				fail = true
+			}
+		}
+	}
+
+	// Create failed sub-results and apply per-test-case traces.
+	// For each test-case event, we capture the trace from first event up until the next test-case event.
+	var testEvent *topdown.Event
+	for i, e := range trace {
+		if e.Op == topdown.TestCaseOp {
+			if testEvent != nil {
+				if p, ok := testCaseTerms(testEvent); ok {
+					if f := result.Update(*p, trace[:i]); f {
+						fail = true
+					}
+				}
+			}
+
+			testEvent = e
+		}
+	}
+	if testEvent != nil {
+		if p, ok := testCaseTerms(testEvent); ok {
+			if f := result.Update(*p, trace); f {
+				fail = true
+			}
+		}
+	}
+
+	return fail, result
+}
+
+func testCaseTerms(e *topdown.Event) (*ast.Array, bool) {
+	if e == nil {
+		return nil, false
+	}
+
+	if expr, ok := e.Node.(*ast.Expr); ok {
+		if arr, ok := expr.Operand(0).Value.(*ast.Array); ok {
+			return arr, true
+		}
+	}
+
+	return nil, false
+}
+
+func subResult(n string, v any) *SubResult {
+	if v == nil {
+		return &SubResult{}
+	}
+
+	switch x := v.(type) {
+	case map[string]any:
+		fail, srs := subResults(x, nil)
+		return &SubResult{
+			Name:       n,
+			Fail:       fail,
+			SubResults: srs,
+		}
+	case bool:
+		return &SubResult{
+			Name: n,
+			Fail: !x,
+		}
+	default:
+		return &SubResult{
+			Name: n,
+			Fail: true,
+		}
+	}
+}
+
 func (r *Runner) runBenchmark(ctx context.Context, txn storage.Transaction, mod *ast.Module, rule *ast.Rule, options BenchmarkOptions) (*Result, bool) {
+	_, rf := ruleName(rule.Head)
+
 	tr := &Result{
 		Location: rule.Loc(),
 		Package:  mod.Package.Path.String(),
-		Name:     rule.Head.Ref().String(), // TODO(sr): test
+		Name:     rf.String(), // TODO(sr): test
 	}

 	var stop bool
@@ -551,14 +1040,23 @@ func (r *Runner) runBenchmark(ctx context.Context, txn storage.Transaction, mod
 		b.ResetTimer()

 		for range b.N {
+			opts := []rego.EvalOption{
+				rego.EvalTransaction(txn),
+				rego.EvalMetrics(m),
+			}
+
+			var tracer *TestQueryTracer
+			if rule.Head.DocKind() == ast.PartialObjectDoc {
+				tracer = NewTestQueryTracer()
+				opts = append(opts, rego.EvalQueryTracer(tracer))
+			}

 			// Start the timer (might already be started, but that's ok)
 			b.StartTimer()

 			rs, err := pq.Eval(
 				ctx,
-				rego.EvalTransaction(txn),
-				rego.EvalMetrics(m),
+				opts...,
 			)

 			// Stop the timer so we don't count any of the error handling time
@@ -573,6 +1071,8 @@ func (r *Runner) runBenchmark(ctx context.Context, txn storage.Transaction, mod
 			} else if len(rs) == 0 {
 				tr.Fail = true
 				b.Fatal("Expected boolean result, got `undefined`")
+			} else if rule.Head.DocKind() == ast.PartialObjectDoc {
+				tr.Fail, tr.SubResults = subResults(rs[0].Expressions[0].Value, tracer.Events())
 			} else if pass, ok := rs[0].Expressions[0].Value.(bool); !ok || !pass {
 				tr.Fail = true
 				b.Fatal("Expected test to evaluate as true, got false")
@@ -53,6 +53,7 @@ type expectedTestResult struct {
 	wantFail bool
 	// nolint: structcheck // The test doesn't check this value, but should.
 	wantSkip bool
+	cases    map[string]expectedTestResult
 }

 type testRunConfig struct {
@@ -67,12 +68,12 @@ func testRun(t *testing.T, conf testRunConfig) map[string]*ast.Module {
 	files := map[string]string{
 		"/a.rego": `package foo
 			import rego.v1
-
+		
 			allow if { true }
 			`,
 		"/a_test.rego": `package foo
 			import rego.v1
-
+		
 			test_pass if { allow }
 			non_test if { true }
 			test_fail if { not allow }
@@ -87,36 +88,69 @@ func testRun(t *testing.T, conf testRunConfig) map[string]*ast.Module {
 			`,
 		"/b_test.rego": `package bar
 			import rego.v1
-
+		
 			test_duplicate if { true }`,
 		"/c_test.rego": `package baz
 			import rego.v1
-
+		
 			a.b.test_duplicate if { false }
 			a.b.test_duplicate if { true }
 			a.b.test_duplicate if { true }`,
 		// Regression test for issue #5496.
 		"/d_test.rego": `package test
-		import rego.v1
+			import rego.v1
+		
+			a[0] := 1
+			test_pass if { true }`,
+		"/e_test.rego": `package qux
+			import rego.v1

-		a[0] := 1
-		test_pass if { true }`,
+			test_cases_pass[x] if { some x in ["foo", "bar"] }
+			test_cases_fail[x] if { some x in ["foo", "bar"]; false }
+			test_cases_partial_fail[x] if { some x in ["foo", "bar", "baz"]; x != "bar" }
+			test_cases_nested[x][y] if { some x in ["foo", "bar"]; some y in ["do", "re", "mi"]; not f(x, y) }
+			f(x, y) if { x == "foo"; y == "re" }`,
 	}

 	tests := expectedTestResults{
-		{"data.foo", "test_pass"}:                  {false, false, false},
-		{"data.foo", "test_fail"}:                  {false, true, false},
-		{"data.foo", "test_fail_non_bool"}:         {false, true, false},
-		{"data.foo", "test_duplicate"}:             {false, true, false},
-		{"data.foo", "test_duplicate#01"}:          {false, false, false},
-		{"data.foo", "test_duplicate#02"}:          {false, false, false},
-		{"data.foo", "test_err"}:                   {true, false, false},
-		{"data.foo", "todo_test_skip"}:             {false, false, true},
-		{"data.bar", "test_duplicate"}:             {false, false, false},
-		{"data.baz", "a.b.test_duplicate"}:         {false, true, false},
-		{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false},
-		{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false},
-		{"data.test", "test_pass"}:                 {false, false, false},
+		{"data.foo", "test_pass"}:                  {false, false, false, nil},
+		{"data.foo", "test_fail"}:                  {false, true, false, nil},
+		{"data.foo", "test_fail_non_bool"}:         {false, true, false, nil},
+		{"data.foo", "test_duplicate"}:             {false, true, false, nil},
+		{"data.foo", "test_duplicate#01"}:          {false, false, false, nil},
+		{"data.foo", "test_duplicate#02"}:          {false, false, false, nil},
+		{"data.foo", "test_err"}:                   {true, false, false, nil},
+		{"data.foo", "todo_test_skip"}:             {false, false, true, nil},
+		{"data.bar", "test_duplicate"}:             {false, false, false, nil},
+		{"data.baz", "a.b.test_duplicate"}:         {false, true, false, nil},
+		{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false, nil},
+		{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false, nil},
+		{"data.test", "test_pass"}:                 {false, false, false, nil},
+		{"data.qux", "test_cases_pass"}: {false, false, false, map[string]expectedTestResult{
+			"foo": {false, false, false, nil},
+			"bar": {false, false, false, nil},
+		}},
+		{"data.qux", "test_cases_fail"}: {false, true, false, map[string]expectedTestResult{
+			"foo": {false, true, false, nil},
+			"bar": {false, true, false, nil},
+		}},
+		{"data.qux", "test_cases_partial_fail"}: {false, true, false, map[string]expectedTestResult{
+			"foo": {false, false, false, nil},
+			"bar": {false, true, false, nil},
+			"baz": {false, false, false, nil},
+		}},
+		{"data.qux", "test_cases_nested"}: {false, true, false, map[string]expectedTestResult{
+			"foo": {false, true, false, map[string]expectedTestResult{
+				"do": {false, false, false, nil},
+				"re": {false, true, false, nil},
+				"mi": {false, false, false, nil},
+			}},
+			"bar": {false, false, false, map[string]expectedTestResult{
+				"do": {false, false, false, nil},
+				"re": {false, false, false, nil},
+				"mi": {false, false, false, nil},
+			}},
+		}},
 	}

 	var modules map[string]*ast.Module
@@ -169,22 +203,27 @@ func doTestRunWithTmpDir(t *testing.T, dir string, conf testRunConfig) ([]*teste
 func validateTestResults(t *testing.T, tests expectedTestResults, rs []*tester.Result, conf testRunConfig) {
 	t.Helper()
 	seen := map[[2]string]struct{}{}
-	for i := range rs {
-		k := [2]string{rs[i].Package, rs[i].Name}
+	for _, r := range rs {
+		k := [2]string{r.Package, r.Name}
 		seen[k] = struct{}{}
 		exp, ok := tests[k]
 		if !ok {
 			t.Errorf("Unexpected result for %v", k)
-		} else if exp.wantErr != (rs[i].Error != nil) || exp.wantFail != rs[i].Fail {
-			t.Errorf("Expected %+v for %v but got: %v", exp, k, rs[i])
+			continue
+		} else if exp.wantErr != (r.Error != nil) || exp.wantFail != r.Fail {
+			t.Errorf("Expected %+v for %v but got: %v", exp, k, r)
 		} else {
 			// Test passed
-			if conf.bench && rs[i].BenchmarkResult == nil {
+			if conf.bench && r.BenchmarkResult == nil {
 				t.Errorf("Expected BenchmarkResult for test %v, got nil", k)
-			} else if !conf.bench && rs[i].BenchmarkResult != nil {
+			} else if !conf.bench && r.BenchmarkResult != nil {
 				t.Errorf("Unexpected BenchmarkResult for test %v, expected nil", k)
 			}
 		}
+
+		if exp.cases != nil {
+			validateSubTestResults(t, exp.cases, r.SubResults)
+		}
 	}
 	for k := range tests {
 		if _, ok := seen[k]; !ok {
@@ -193,6 +232,32 @@ func validateTestResults(t *testing.T, tests expectedTestResults, rs []*tester.R
 	}
 }

+func validateSubTestResults(t *testing.T, tests map[string]expectedTestResult, srs tester.SubResultMap) {
+	t.Helper()
+	seen := map[string]struct{}{}
+	for k, exp := range tests {
+		seen[k] = struct{}{}
+		sr, ok := srs[k]
+		if !ok {
+			t.Errorf("Expected sub-result for %v", k)
+			continue
+		}
+
+		if exp.wantFail != sr.Fail {
+			t.Errorf("Expected %+v for %v but got: %v", exp, k, sr)
+		}
+	}
+	for k, v := range srs {
+		if _, ok := seen[k]; !ok {
+			t.Errorf("Expected sub-result for %v", k)
+		}
+
+		if v.SubResults != nil {
+			validateSubTestResults(t, tests[k].cases, v.SubResults)
+		}
+	}
+}
+
 func TestRunWithFilterRegex(t *testing.T) {
 	files := map[string]string{
 		"/a.rego": `package foo
@@ -215,6 +280,9 @@ func TestRunWithFilterRegex(t *testing.T) {
 			test_duplicate if { true }
 			todo_test_skip if { true }
 			todo_test_skip_too if { false }
+			test_cases[x][y] if { x := "foo"; y := "bar" }
+			test_duplicate.foo[y] if { x := "foo"; y := "bar" }
+			test_duplicate[x][y] if { x := "foo"; y := "bar" }
 			`,
 		"/b_test.rego": `package bar
 			import rego.v1
@@ -237,38 +305,44 @@ func TestRunWithFilterRegex(t *testing.T) {
 			note:  "all tests match",
 			regex: ".*",
 			tests: expectedTestResults{
-				{"data.foo", "test_pass"}:                  {false, false, false},
-				{"data.foo", "test_fail"}:                  {false, true, false},
-				{"data.foo", "test_fail_non_bool"}:         {false, true, false},
-				{"data.foo", "test_duplicate"}:             {false, true, false},
-				{"data.foo", "test_duplicate#01"}:          {false, false, false},
-				{"data.foo", "test_duplicate#02"}:          {false, false, false},
-				{"data.foo", "test_err"}:                   {true, false, false},
-				{"data.foo", "todo_test_skip"}:             {false, false, true},
-				{"data.foo", "todo_test_skip_too"}:         {false, false, true},
-				{"data.bar", "test_duplicate"}:             {false, false, false},
-				{"data.baz", "a.b.test_duplicate"}:         {false, true, false},
-				{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false},
-				{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false},
+				{"data.foo", "test_pass"}:                  {false, false, false, nil},
+				{"data.foo", "test_fail"}:                  {false, true, false, nil},
+				{"data.foo", "test_fail_non_bool"}:         {false, true, false, nil},
+				{"data.foo", "test_duplicate"}:             {false, true, false, nil},
+				{"data.foo", "test_duplicate#01"}:          {false, false, false, nil},
+				{"data.foo", "test_duplicate#02"}:          {false, false, false, nil},
+				{"data.foo", "test_err"}:                   {true, false, false, nil},
+				{"data.foo", "todo_test_skip"}:             {false, false, true, nil},
+				{"data.foo", "todo_test_skip_too"}:         {false, false, true, nil},
+				{"data.foo", "test_cases"}:                 {false, false, false, nil},
+				{"data.foo", "test_duplicate#03"}:          {false, false, false, nil},
+				{"data.foo", "test_duplicate#04"}:          {false, false, false, nil},
+				{"data.bar", "test_duplicate"}:             {false, false, false, nil},
+				{"data.baz", "a.b.test_duplicate"}:         {false, true, false, nil},
+				{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false, nil},
+				{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false, nil},
 			},
 		},
 		{
 			note:  "no filter",
 			regex: "",
 			tests: expectedTestResults{
-				{"data.foo", "test_pass"}:                  {false, false, false},
-				{"data.foo", "test_fail"}:                  {false, true, false},
-				{"data.foo", "test_fail_non_bool"}:         {false, true, false},
-				{"data.foo", "test_duplicate"}:             {false, true, false},
-				{"data.foo", "test_duplicate#01"}:          {false, false, false},
-				{"data.foo", "test_duplicate#02"}:          {false, false, false},
-				{"data.foo", "test_err"}:                   {true, false, false},
-				{"data.foo", "todo_test_skip"}:             {false, false, true},
-				{"data.foo", "todo_test_skip_too"}:         {false, false, true},
-				{"data.bar", "test_duplicate"}:             {false, false, false},
-				{"data.baz", "a.b.test_duplicate"}:         {false, true, false},
-				{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false},
-				{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false},
+				{"data.foo", "test_pass"}:                  {false, false, false, nil},
+				{"data.foo", "test_fail"}:                  {false, true, false, nil},
+				{"data.foo", "test_fail_non_bool"}:         {false, true, false, nil},
+				{"data.foo", "test_duplicate"}:             {false, true, false, nil},
+				{"data.foo", "test_duplicate#01"}:          {false, false, false, nil},
+				{"data.foo", "test_duplicate#02"}:          {false, false, false, nil},
+				{"data.foo", "test_err"}:                   {true, false, false, nil},
+				{"data.foo", "todo_test_skip"}:             {false, false, true, nil},
+				{"data.foo", "todo_test_skip_too"}:         {false, false, true, nil},
+				{"data.foo", "test_cases"}:                 {false, false, false, nil},
+				{"data.foo", "test_duplicate#03"}:          {false, false, false, nil},
+				{"data.foo", "test_duplicate#04"}:          {false, false, false, nil},
+				{"data.bar", "test_duplicate"}:             {false, false, false, nil},
+				{"data.baz", "a.b.test_duplicate"}:         {false, true, false, nil},
+				{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false, nil},
+				{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false, nil},
 			},
 		},
 		{
@@ -280,67 +354,74 @@ func TestRunWithFilterRegex(t *testing.T) {
 			note:  "single package name",
 			regex: "bar",
 			tests: expectedTestResults{
-				{"data.bar", "test_duplicate"}: {false, false, false},
+				{"data.bar", "test_duplicate"}: {false, false, false, nil},
 			},
 		},
 		{
 			note:  "single package explicit",
 			regex: "data.bar.test_duplicate",
 			tests: expectedTestResults{
-				{"data.bar", "test_duplicate"}: {false, false, false},
+				{"data.bar", "test_duplicate"}: {false, false, false, nil},
 			},
 		},
 		{
-			note:  "single test ",
+			note:  "single test",
 			regex: "test_pass",
 			tests: expectedTestResults{
-				{"data.foo", "test_pass"}: {false, false, false},
+				{"data.foo", "test_pass"}: {false, false, false, nil},
 			},
 		},
 		{
 			note:  "single test explicit",
 			regex: "data.foo.test_pass",
 			tests: expectedTestResults{
-				{"data.foo", "test_pass"}: {false, false, false},
+				{"data.foo", "test_pass"}: {false, false, false, nil},
 			},
 		},
 		{
 			note:  "single test skipped explicit",
 			regex: "data.foo.todo_test_skip_too",
 			tests: expectedTestResults{
-				{"data.foo", "todo_test_skip_too"}: {false, false, true},
+				{"data.foo", "todo_test_skip_too"}: {false, false, true, nil},
 			},
 		},
 		{
 			note:  "wildcards",
 			regex: "^.*foo.*_fail.*$",
 			tests: expectedTestResults{
-				{"data.foo", "test_fail"}:          {false, true, false},
-				{"data.foo", "test_fail_non_bool"}: {false, true, false},
+				{"data.foo", "test_fail"}:          {false, true, false, nil},
+				{"data.foo", "test_fail_non_bool"}: {false, true, false, nil},
 			},
 		},
 		{
 			note:  "mixed",
 			regex: "(bar|data.foo.test_pass)",
 			tests: expectedTestResults{
-				{"data.foo", "test_pass"}:      {false, false, false},
-				{"data.bar", "test_duplicate"}: {false, false, false},
+				{"data.foo", "test_pass"}:      {false, false, false, nil},
+				{"data.bar", "test_duplicate"}: {false, false, false, nil},
 			},
 		},
 		{
 			note:  "case insensitive",
 			regex: "(?i)DATA.BAR",
 			tests: expectedTestResults{
-				{"data.bar", "test_duplicate"}: {false, false, false},
+				{"data.bar", "test_duplicate"}: {false, false, false, nil},
 			},
 		},
 		{
 			note:  "matching ref rule halfways",
 			regex: "data.baz.a",
 			tests: expectedTestResults{
-				{"data.baz", "a.b.test_duplicate"}:         {false, true, false},
-				{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false},
-				{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false},
+				{"data.baz", "a.b.test_duplicate"}:         {false, true, false, nil},
+				{"data.baz", "a.b[\"test_duplicate#01\"]"}: {false, false, false, nil},
+				{"data.baz", "a.b[\"test_duplicate#02\"]"}: {false, false, false, nil},
+			},
+		},
+		{
+			note:  "matching sub-test rule",
+			regex: "data.foo.test_cases",
+			tests: expectedTestResults{
+				{"data.foo", "test_cases"}: {false, false, false, nil},
 			},
 		},
 	}
@@ -0,0 +1,28 @@
+// Copyright 2025 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package tester
+
+import "github.com/open-policy-agent/opa/v1/topdown"
+
+type TestQueryTracer struct {
+	topdown.BufferTracer
+}
+
+func NewTestQueryTracer() *TestQueryTracer {
+	return &TestQueryTracer{}
+}
+
+func (t *TestQueryTracer) TraceEvent(e topdown.Event) {
+	if e.Op == topdown.TestCaseOp {
+		t.BufferTracer.TraceEvent(e)
+	}
+}
+
+func (t *TestQueryTracer) Events() []*topdown.Event {
+	if t == nil {
+		return nil
+	}
+	return t.BufferTracer
+}
@@ -0,0 +1,30 @@
+// Copyright 2025 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package topdown
+
+import "github.com/open-policy-agent/opa/v1/ast"
+
+const TestCaseOp Op = "TestCase"
+
+func builtinTestCase(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
+	e := &Event{
+		Op:      TestCaseOp,
+		QueryID: bctx.QueryID,
+		Node: ast.NewExpr([]*ast.Term{
+			ast.NewTerm(ast.InternalTestCase.Ref()),
+			ast.NewTerm(operands[0].Value),
+		}),
+	}
+
+	for _, tracer := range bctx.QueryTracers {
+		tracer.TraceEvent(*e)
+	}
+
+	return iter(ast.BooleanTerm(true))
+}
+
+func init() {
+	RegisterBuiltinFunc(ast.InternalTestCase.Name, builtinTestCase)
+}