rego-v1: Future-proofing internal tests to be 1.0 compatible (#7020)

Signed-off-by: Johan Fylling <johan.dev@fylling.se>

internal/compiler/utils_test.go

@@ -17,39 +17,37 @@ func TestVerifyAuthorizationPolicySchema(t *testing.T) {
 	module1 := `
 	package policy
 
-    import future.keywords
-
 	default allow := false
 
-	allow {
+	allow if {
 	   input.identity = "foo"
 	}
 
-	allow {
+	allow if {
 	   input.client_certificates[0] = {"foo": "bar"}
 	}
 
-	allow {
+	allow if {
 	   input.method = "GET"
 	}
 
-	allow {
+	allow if {
 	   input.path = ["foo", "bar"]
 	}
 
-    allow {
+    allow if {
 	   "foo" in input.path
 	}
 
-	allow {
+	allow if {
 	   input.params = {"foo": "bar"}
 	}
 
-	allow {
+	allow if {
 	   input.headers = {"foo": "bar"}
 	}
 
-	allow {
+	allow if {
 	   input.body.input.stock = "ACME"
 	}`
 
@@ -58,11 +56,11 @@ func TestVerifyAuthorizationPolicySchema(t *testing.T) {
 
 	default allow := false
 
-	allow {
+	allow if {
 	   input.identty = "foo"
 	}
 
-	allow {
+	allow if {
 	   input.path = "foo"
 	}`
 
@@ -71,7 +69,7 @@ func TestVerifyAuthorizationPolicySchema(t *testing.T) {
 
 	default allow := false
 
-	allow {
+	allow if {
 	   input.path = [1, 2, 3]
 	}`
 
@@ -80,7 +78,7 @@ func TestVerifyAuthorizationPolicySchema(t *testing.T) {
 
     default allow := false
 
-    allow {
+    allow if {
        input.client_certificates[0] = "foo"
     }`
 
@@ -102,7 +100,8 @@ func TestVerifyAuthorizationPolicySchema(t *testing.T) {
 			modules := map[string]*ast.Module{}
 
 			for i, module := range tc.modules {
-				mod, err := ast.ParseModule(fmt.Sprintf("test%d.rego", i+1), module)
+				mod, err := ast.ParseModuleWithOpts(fmt.Sprintf("test%d.rego", i+1), module,
+					ast.ParserOptions{AllFutureKeywords: true})
 				if err != nil {
 					t.Fatal(err)
 				}

internal/oracle/oracle_test.go

@@ -33,35 +33,39 @@ f(input)`,
 		{
 			note: "no matching node",
 			buffer: `package test
+import rego.v1
 
-p { q }
+p if { q }
 
 q = true`,
-			pos: 100,
+			pos: 118,
 			exp: ErrNoMatchFound,
 		},
 		{
 			note: "no good match - literal",
 			buffer: `package test
+import rego.v1
 
-p { q > 1 }`,
-			pos: 22, // this points at the number '1'
+p if { q > 1 }`,
+			pos: 40, // this points at the number '1'
 			exp: ErrNoDefinitionFound,
 		},
 		{
 			note: "no good match - rule name",
 			buffer: `package test
+import rego.v1
 
-p { q > 1 }`,
-			pos: 14, // this points at the rule 'p'
+p if { q > 1 }`,
+			pos: 32, // this points at the rule 'p'
 			exp: ErrNoDefinitionFound,
 		},
 		{
 			note: "no good match - rule whitespace",
 			buffer: `package test
+import rego.v1
 
-p { q > 1 }`,
-			pos: 21, // this points at the whitespace after '>'
+p if { q > 1 }`,
+			pos: 39, // this points at the whitespace after '>'
 			exp: ErrNoDefinitionFound,
 		},
 		{
@@ -105,10 +109,11 @@ func TestOracleFindDefinition(t *testing.T) {
 
 	const aBufferModule = `package test
 
+import rego.v1
 import data.foo.s
 import data.foo.bar as t
 
-p {
+p if {
     q
     [r]
     s[t]
@@ -118,39 +123,41 @@ r = true
 q = true`
 
 	const aSecondBufferModule = `package test
+import rego.v1
 
-p {
+p if {
 	q
 }`
 
 	const aThirdBufferModule = `package test
+import rego.v1
 
-f(x) {
+f(x) if {
 	input.foo[x]
 }
 
-u {
+u if {
 	some x
 	x = 1
 }
 
-v {
+v if {
 	x := 1
 	x < 10
 }
 
-w {
+w if {
 	y[i]
 	i > 1
 }
 
-m {
+m if {
 	[i, j] = [1, 2]
 	j > i
 }
 
 x = "deadbeef"
-y[1]
+y contains 1
 `
 
 	const fooModule = `package foo
@@ -160,7 +167,9 @@ bar = 7`
 
 	// NOTE(sr): Early ref rewriting adds an expression to the rule body for `x.y`
 	const varInRuleRefModule = `package foo
-q[x.y] = 10 {
+import rego.v1
+
+q[x.y] = 10 if {
 	x := input
 	some z
 	z = 1
@@ -176,10 +185,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aBufferModule,
 			},
-			pos: 66,
+			pos: 84,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    13,
+				Row:    14,
 				Col:    1,
 				Offset: 97,
 				Text:   []byte("q = true"),
@@ -190,10 +199,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aBufferModule,
 			},
-			pos: 73,
+			pos: 91,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    12,
+				Row:    13,
 				Col:    1,
 				Offset: 88,
 				Text:   []byte("r = true"),
@@ -205,7 +214,7 @@ q[x.y] = 10 {
 				"buffer.rego": aBufferModule,
 				"foo.rego":    fooModule,
 			},
-			pos: 80,
+			pos: 98,
 			exp: &ast.Location{
 				File:   "foo.rego",
 				Row:    3,
@@ -220,7 +229,7 @@ q[x.y] = 10 {
 				"buffer.rego": aBufferModule,
 				"foo.rego":    fooModule,
 			},
-			pos: 81, // this refers to the '[' character following 's'--this exercises the case where position does not refer to a symbol
+			pos: 99, // this refers to the '[' character following 's'--this exercises the case where position does not refer to a symbol
 			exp: &ast.Location{
 				File:   "foo.rego",
 				Row:    3,
@@ -234,10 +243,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aBufferModule,
 			},
-			pos: 80,
+			pos: 98,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    3,
+				Row:    4,
 				Col:    8,
 				Offset: 21,
 				Text:   []byte("data.foo.s"),
@@ -249,7 +258,7 @@ q[x.y] = 10 {
 				"buffer.rego": aBufferModule,
 				"foo.rego":    fooModule,
 			},
-			pos: 82,
+			pos: 100,
 			exp: &ast.Location{
 				File:   "foo.rego",
 				Row:    4,
@@ -263,10 +272,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aBufferModule,
 			},
-			pos: 82,
+			pos: 100,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    4,
+				Row:    5,
 				Col:    8,
 				Offset: 39,
 				Text:   []byte("data.foo.bar"),
@@ -278,10 +287,10 @@ q[x.y] = 10 {
 				"buffer.rego": aSecondBufferModule, // use a different module that references q in main buffer module used above
 				"test.rego":   aBufferModule,
 			},
-			pos: 19,
+			pos: 37,
 			exp: &ast.Location{
 				File:   "test.rego",
-				Row:    13,
+				Row:    14,
 				Col:    1,
 				Offset: 97,
 				Text:   []byte("q = true"),
@@ -292,10 +301,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aThirdBufferModule,
 			},
-			pos: 32,
+			pos: 50,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    3,
+				Row:    4,
 				Col:    3,
 				Offset: 16,
 				Text:   []byte("x"),
@@ -306,10 +315,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aThirdBufferModule,
 			},
-			pos: 51,
+			pos: 72,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    8,
+				Row:    9,
 				Col:    7,
 				Offset: 48,
 				Text:   []byte("x"),
@@ -320,10 +329,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aThirdBufferModule,
 			},
-			pos: 73,
+			pos: 97,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    13,
+				Row:    14,
 				Col:    2,
 				Offset: 65,
 				Text:   []byte("x"),
@@ -334,10 +343,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aThirdBufferModule,
 			},
-			pos: 94,
+			pos: 121,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    18,
+				Row:    19,
 				Col:    4,
 				Offset: 90,
 				Text:   []byte("i"),
@@ -348,10 +357,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": aThirdBufferModule,
 			},
-			pos: 129,
+			pos: 159,
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    23,
+				Row:    24,
 				Col:    3,
 				Offset: 109,
 				Text:   []byte("i"),
@@ -362,10 +371,10 @@ q[x.y] = 10 {
 			modules: map[string]string{
 				"buffer.rego": varInRuleRefModule,
 			},
-			pos: 47, // "z" in "z = 1"
+			pos: 66, // "z" in "z = 1"
 			exp: &ast.Location{
 				File:   "buffer.rego",
-				Row:    4,
+				Row:    6,
 				Col:    7,
 				Offset: 44,
 				Text:   []byte("z"),
@@ -420,8 +429,9 @@ q[x.y] = 10 {
 
 func TestFindContainingNodeStack(t *testing.T) {
 	const trivial = `package test
+import rego.v1
 
-p {
+p if {
     q
     r
 }
@@ -432,8 +442,8 @@ q = true`
 	module := ast.MustParseModule(trivial)
 	module.Package.Location = nil // unset the package location to test nil tolerance
 
-	// offset 28 is the first 'r' variable
-	result := findContainingNodeStack(module, 28)
+	// offset 46 is the first 'r' variable
+	result := findContainingNodeStack(module, 46)
 
 	exp := []*ast.Location{
 		module.Rules[0].Loc(),
@@ -454,7 +464,7 @@ q = true`
 
 	// Exercise special case for bodies.
 	module.Rules[0].Body[1].Location = nil
-	result = findContainingNodeStack(module, 28)
+	result = findContainingNodeStack(module, 46)
 
 	exp = []*ast.Location{
 		module.Rules[0].Loc(),

internal/planner/planner_test.go

@@ -89,8 +89,8 @@ func TestPlannerHelloWorld(t *testing.T) {
 			queries: []string{"data.test.p = x"},
 			modules: []string{`
 				package test
-				p = x { x = 1 }
-				p = y { y = 2 }
+				p = x if { x = 1 }
+				p = y if { y = 2 }
 			`},
 		},
 		{
@@ -98,7 +98,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 			queries: []string{"data.test.p = 10"},
 			modules: []string{`
 				package test
-				p = x { x = 10 }
+				p = x if { x = 10 }
 			`},
 		},
 		{
@@ -106,7 +106,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 			queries: []string{"data.test.f([1,x])"},
 			modules: []string{`
 				package test
-				f([a, b]) {
+				f([a, b]) if {
 					a = b
 				}
 			`},
@@ -116,9 +116,9 @@ func TestPlannerHelloWorld(t *testing.T) {
 			queries: []string{"data.test.p = 1"},
 			modules: []string{`
 				package test
-				p = 0 {
+				p = 0 if {
 					false
-				} else = 1 {
+				} else = 1 if {
 					true
 				}
 			`},
@@ -128,8 +128,8 @@ func TestPlannerHelloWorld(t *testing.T) {
 			queries: []string{"data.test.p = {1,2}"},
 			modules: []string{`
 				package test
-				p[1]
-				p[2]
+				p contains 1
+				p contains 2
 			`},
 		},
 		{
@@ -147,7 +147,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{`
 				package test
 				p["a"] = 1
-				p[v] = 2 { v := "b" }
+				p[v] = 2 if { v := "b" }
 			`},
 		},
 		{
@@ -156,7 +156,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{`
 				package test
 				p.q.r["a"] = 1
-				p.q[v] = 2 { v := "b" }
+				p.q[v] = 2 if { v := "b" }
 			`},
 		},
 		{
@@ -165,9 +165,9 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{`
 				package test
 				p.q["a"] = 1
-				p.q[v] = 2 { v := "b" }
+				p.q[v] = 2 if { v := "b" }
 				p.r["c"] = 3
-				p.r[v] = 4 { v := "d" }
+				p.r[v] = 4 if { v := "d" }
 			`},
 		},
 		{
@@ -176,7 +176,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{`
 				package test
 				p.q["a"] = 1
-				p.q[v] = x { l1 := ["b", "c", "d"]; l2 := ["foo", "bar"]; l3 := [2, 3]; v := l1[_]; x := l2[_]; z := l3[_] }
+				p.q[v] = x if { l1 := ["b", "c", "d"]; l2 := ["foo", "bar"]; l3 := [2, 3]; v := l1[_]; x := l2[_]; z := l3[_] }
 			`},
 		},
 		{
@@ -187,7 +187,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 				p.q["a"] = 1
 				p.q.b.s.baz = 2
 				p.q.b.s.foo.c = 3
-				p.q[r].s[t].u = v { x := ["foo", "bar"]; r := "b"; t := x[v]}
+				p.q[r].s[t].u = v if { x := ["foo", "bar"]; r := "b"; t := x[v]}
 			`},
 		},
 		{
@@ -195,7 +195,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 			queries: []string{`data.test.p`},
 			modules: []string{`
 				package test
-				p { xs = [1]; every k, v in xs { k < v } }
+				p if { xs = [1]; every k, v in xs { k < v } }
 			`},
 		},
 		{
@@ -205,7 +205,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 				package test
 
 				p = 1
-				q = 2 { false }
+				q = 2 if { false }
 			`},
 		},
 		{
@@ -233,7 +233,7 @@ func TestPlannerHelloWorld(t *testing.T) {
 					package test.a
 
 					p = 1
-					q = 2 { false }
+					q = 2 if { false }
 					r = 3
 				`,
 				`
@@ -291,10 +291,10 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{
 				`package p
 
-				q {
+				q if {
 					false
 				}
-				else = true {
+				else = true if {
 					true
 				}
 				q = false
@@ -307,16 +307,16 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{
 				`package p
 
-				q {
+				q if {
 					false
 				}
-				else = false {
+				else = false if {
 					true
 				}
-				q {
+				q if {
 					false
 				}
-				else = true {
+				else = true if {
 					true
 				}`,
 			},
@@ -327,11 +327,11 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{
 				`package p
 
-				p(a) = y {
+				p(a) = y if {
 				  y = a[_]
 				}
 
-				r = y {
+				r = y if {
 				  data.p.p([1, 2, 3], y)
 				}
 				`,
@@ -343,14 +343,14 @@ func TestPlannerHelloWorld(t *testing.T) {
 			modules: []string{
 				`package p
 
-				p(1, a) = y {
+				p(1, a) = y if {
 					y = a
 				}
-				p(x, y) = z {
+				p(x, y) = z if {
 					z = x
 				}
 
-				r = y {
+				r = y if {
 					data.p.p(1, 0, y)
 				}
 				`,
@@ -493,7 +493,7 @@ func TestPlannerLocations(t *testing.T) {
 			queries: []string{"data.test.p = 10"},
 			modules: []string{`
 package test
-p = x {
+p = x if {
   1 > 0
   x = 10
   true
@@ -511,12 +511,12 @@ p = x {
 			queries: []string{"data.test.p = {1,2}"},
 			modules: []string{`
 package test
-p[1]
-p[2]
+p contains 1
+p contains 2
 			`},
 			exps: map[ir.Stmt]string{
-				&ir.MakeSetStmt{}:     "module-0.rego:3:1: p[1]",
-				&ir.ReturnLocalStmt{}: "module-0.rego:3:1: p[1]",
+				&ir.MakeSetStmt{}:     "module-0.rego:3:1: p contains 1",
+				&ir.ReturnLocalStmt{}: "module-0.rego:3:1: p contains 1",
 			},
 			where: funcs,
 		},
@@ -525,13 +525,13 @@ p[2]
 			queries: []string{"data.test.p = {1,2}"},
 			modules: []string{`
 package test
-p[1] {
+p contains 1 if {
   1 > 2
 }
 			`},
 			exps: map[ir.Stmt]string{
 				&ir.CallStmt{}:   "module-0.rego:4:3: 1 > 2",
-				&ir.SetAddStmt{}: "module-0.rego:3:1: p[1]",
+				&ir.SetAddStmt{}: "module-0.rego:3:1: p contains 1",
 			},
 			where: funcs,
 		},
@@ -540,7 +540,7 @@ p[1] {
 			queries: []string{`data.test.p = {"a": 1, "b": 2}`},
 			modules: []string{`
 package test
-p["a"] = 1 {
+p["a"] = 1 if {
   false
 }
 			`},
@@ -556,7 +556,7 @@ p["a"] = 1 {
 			modules: []string{`
 package test
 default p = {"foo": "bar"}
-p = x {
+p = x if {
   x := {"baz": "quz"}
 }
 			`},
@@ -625,7 +625,7 @@ a = { "a", 10 }`},
 			queries: []string{`data`},
 			modules: []string{`package test
 p = 1
-q = 2 {
+q = 2 if {
   false
 }`},
 			exps: map[ir.Stmt]string{
@@ -656,7 +656,7 @@ a = true`},
 			note:    "non-ground ref in policy",
 			queries: []string{`data.test.a = x`},
 			modules: []string{`package test
-a {
+a if {
   data.test1[_].y = "z"
 }`},
 			exps: map[ir.Stmt]string{
@@ -671,7 +671,7 @@ a {
 			note:    "CallDynamicStmt optimization",
 			queries: []string{`x := "a"; data.test[x] = y`},
 			modules: []string{`package test
-a {
+a if {
   true
 }`},
 			exps: map[ir.Stmt]string{
@@ -689,7 +689,7 @@ a {
 			modules := make([]*ast.Module, len(tc.modules))
 			for i := range modules {
 				file := fmt.Sprintf("module-%d.rego", i)
-				m, err := ast.ParseModule(file, tc.modules[i])
+				m, err := ast.ParseModuleWithOpts(file, tc.modules[i], ast.ParserOptions{AllFutureKeywords: true})
 				if err != nil {
 					t.Fatal(err)
 				}
@@ -1084,7 +1084,7 @@ func TestPlannerCallDynamic(t *testing.T) {
 			note:    "CallDynamicStmt optimization",
 			queries: []string{`x := "a"; data.test[x] = y`},
 			modules: []string{`package test
-a { true }`},
+a if { true }`},
 			path: []interface{}{"g0", "test", 2},
 			extras: []func(interface{}) error{
 				findFunc("g0.data.test.a", "g0.test.a"),
@@ -1094,7 +1094,7 @@ a { true }`},
 			note:    "simple single-val ref head",
 			queries: []string{`x := "a"; data.test.a[x].c = y`},
 			modules: []string{`package test
-a.b.c = 1 { true }`},
+a.b.c = 1 if { true }`},
 			path: []interface{}{"g0", "test", "a", 2, "c"},
 			extras: []func(interface{}) error{
 				findFunc("g0.data.test.a.b.c", "g0.test.a.b.c"),
@@ -1104,8 +1104,8 @@ a.b.c = 1 { true }`},
 			note:    "two single-val ref heads, string+var",
 			queries: []string{`x := "a"; data.test.a[x] = y`},
 			modules: []string{`package test
-a.b.c = 1 { true }
-a.b[t] = 2 { t := input }`},
+a.b.c = 1 if { true }
+a.b[t] = 2 if { t := input }`},
 			path: []interface{}{"g0", "test", "a", 2},
 			extras: []func(interface{}) error{
 				findFunc("g0.data.test.a.b", "g0.test.a.b"),
@@ -1115,8 +1115,8 @@ a.b[t] = 2 { t := input }`},
 			note:    "two single-val ref heads, number+var",
 			queries: []string{`x := "a"; data.test.a[x] = y`},
 			modules: []string{`package test
-a.b[1] = 1 { true }
-a.b[t] = 2 { t := input }`},
+a.b[1] = 1 if { true }
+a.b[t] = 2 if { t := input }`},
 			path: []interface{}{"g0", "test", "a", 2},
 			extras: []func(interface{}) error{
 				findFunc("g0.data.test.a.b", "g0.test.a.b"),
@@ -1126,7 +1126,7 @@ a.b[t] = 2 { t := input }`},
 			note:    "one single-val ref head, number",
 			queries: []string{`x := "a"; data.test.a[x] = y`},
 			modules: []string{`package test
-a.b[1] = 1 { true }`},
+a.b[1] = 1 if { true }`},
 			path: []interface{}{"g0", "test", "a", 2},
 			extras: []func(interface{}) error{
 				findFunc("g0.data.test.a.b", "g0.test.a.b"),
@@ -1143,7 +1143,7 @@ a.b[1] = 1 { true }`},
 			modules := make([]*ast.Module, len(tc.modules))
 			for i := range modules {
 				file := fmt.Sprintf("module-%d.rego", i)
-				m, err := ast.ParseModule(file, tc.modules[i])
+				m, err := ast.ParseModuleWithOpts(file, tc.modules[i], ast.ParserOptions{AllFutureKeywords: true})
 				if err != nil {
 					t.Fatal(err)
 				}

internal/presentation/presentation_test.go

@@ -144,8 +144,9 @@ func TestOutputJSONErrorStructuredStorageErr(t *testing.T) {
 func TestOutputJSONErrorStructuredTopdownErr(t *testing.T) {
 	mod := `
 		package test
+		import rego.v1
 
-		p(x) = y {
+		p(x) = y if {
 			y = x[_]
 		}
 
@@ -164,7 +165,7 @@ func TestOutputJSONErrorStructuredTopdownErr(t *testing.T) {
       "code": "eval_conflict_error",
       "location": {
         "file": "test.rego",
-        "row": 4,
+        "row": 5,
         "col": 3
       }
     }
@@ -277,8 +278,9 @@ func TestOutputJSONErrorStructuredASTErrList(t *testing.T) {
 	c.Compile(map[string]*ast.Module{
 		"error.rego": ast.MustParseModule(`
 package test
+import rego.v1
 
-q {
+q if {
 	bad[reference]
 }
 `)})
@@ -292,7 +294,7 @@ q {
       "code": "rego_unsafe_var_error",
       "location": {
         "file": "",
-        "row": 5,
+        "row": 6,
         "col": 2
       }
     },
@@ -301,7 +303,7 @@ q {
       "code": "rego_unsafe_var_error",
       "location": {
         "file": "",
-        "row": 5,
+        "row": 6,
         "col": 2
       }
     }
@@ -344,12 +346,13 @@ func TestOutputJSONErrorStructuredLoaderErrList(t *testing.T) {
 func TestOutputJSONErrorStructuredRegoErrList(t *testing.T) {
 	mod := `
 package test
+import rego.v1
 
-p {
+p if {
 	bad_func1()
 }
 
-q {
+q if {
 	bad_func2()
 }
 `
@@ -365,7 +368,7 @@ q {
       "code": "rego_type_error",
       "location": {
         "file": "error.rego",
-        "row": 5,
+        "row": 6,
         "col": 2
       }
     },
@@ -374,7 +377,7 @@ q {
       "code": "rego_type_error",
       "location": {
         "file": "error.rego",
-        "row": 9,
+        "row": 10,
         "col": 2
       }
     }
@@ -397,7 +400,7 @@ func TestSource(t *testing.T) {
 			Support: []*ast.Module{
 				ast.MustParseModule(`
             package test
-            p = 1
+            p := 1
         `),
 			},
 		},
@@ -413,7 +416,7 @@ b = 2
 # Module 1
 package test
 
-p = 1
+p := 1
 `
 
 	if buf.String() != exp {