1// Copyright 2016 The OPA Authors.  All rights reserved.
2// Use of this source code is governed by an Apache2
3// license that can be found in the LICENSE file.
4
5package repl
6
7import (
8	"bytes"
9	"context"
10	"encoding/json"
11	"fmt"
12	"io"
13	"io/ioutil"
14	"os"
15	"path/filepath"
16	"reflect"
17	"sort"
18	"strings"
19	"testing"
20
21	"github.com/open-policy-agent/opa/ast"
22	"github.com/open-policy-agent/opa/storage"
23	"github.com/open-policy-agent/opa/storage/inmem"
24	"github.com/open-policy-agent/opa/util"
25)
26
27func TestFunction(t *testing.T) {
28	store := newTestStore()
29	ctx := context.Background()
30	txn := storage.NewTransactionOrDie(ctx, store, storage.WriteParams)
31
32	mod1 := []byte(`package a.b.c
33
34foo(x) = y {
35	split(x, ".", y)
36}
37
38bar([x, y]) = z {
39	trim(x, y, z)
40}
41`)
42
43	mod2 := []byte(`package a.b.d
44
45baz(_) = y {
46	data.a.b.c.foo("barfoobar.bar", x)
47	data.a.b.c.bar(x, y)
48}`)
49
50	if err := store.UpsertPolicy(ctx, txn, "mod1", mod1); err != nil {
51		panic(err)
52	}
53
54	if err := store.UpsertPolicy(ctx, txn, "mod2", mod2); err != nil {
55		panic(err)
56	}
57
58	if err := store.Commit(ctx, txn); err != nil {
59		panic(err)
60	}
61
62	var buf bytes.Buffer
63	repl := newRepl(store, &buf)
64	repl.OneShot(ctx, "json")
65	repl.OneShot(ctx, "data.a.b.d.baz(null, x)")
66	exp := util.MustUnmarshalJSON([]byte(`[{"x": "foo"}]`))
67	result := util.MustUnmarshalJSON(buf.Bytes())
68	if !reflect.DeepEqual(exp, result) {
69		t.Fatalf("expected data.a.b.d.baz(x) to be %v, got %v", exp, result)
70	}
71
72	err := repl.OneShot(ctx, "p(x) = y { y = x+4 }")
73	if err != nil {
74		t.Fatalf("failed to compile repl function: %v", err)
75	}
76
77	buf.Reset()
78	repl.OneShot(ctx, "data.repl.p(5, y)")
79	exp = util.MustUnmarshalJSON([]byte(`[{"y": 9}]`))
80	result = util.MustUnmarshalJSON(buf.Bytes())
81	if !reflect.DeepEqual(exp, result) {
82		t.Fatalf("expected datrepl.p(x) to be %v, got %v", exp, result)
83	}
84
85	repl.OneShot(ctx, "f(1, x) = y { y = x }")
86	repl.OneShot(ctx, "f(2, x) = y { y = x*2 }")
87
88	buf.Reset()
89	repl.OneShot(ctx, "data.repl.f(1, 2, y)")
90	exp = util.MustUnmarshalJSON([]byte(`[{"y": 2}]`))
91	result = util.MustUnmarshalJSON(buf.Bytes())
92	if !reflect.DeepEqual(exp, result) {
93		t.Fatalf("expected data.repl.f(1, 2, y) to be %v, got %v", exp, result)
94	}
95	buf.Reset()
96	repl.OneShot(ctx, "data.repl.f(2, 2, y)")
97	exp = util.MustUnmarshalJSON([]byte(`[{"y": 4}]`))
98	result = util.MustUnmarshalJSON(buf.Bytes())
99	if !reflect.DeepEqual(exp, result) {
100		t.Fatalf("expected data.repl.f(2, 2, y) to be %v, got %v", exp, result)
101	}
102}
103
104func TestComplete(t *testing.T) {
105	ctx := context.Background()
106	store := newTestStore()
107	txn := storage.NewTransactionOrDie(ctx, store, storage.WriteParams)
108
109	mod1 := []byte(`package a.b.c
110
111p = 1 { true }
112q = 2 { true }
113q = 3 { false }`)
114
115	mod2 := []byte(`package a.b.d
116
117r = 3 { true }`)
118
119	if err := store.UpsertPolicy(ctx, txn, "mod1", mod1); err != nil {
120		panic(err)
121	}
122
123	if err := store.UpsertPolicy(ctx, txn, "mod2", mod2); err != nil {
124		panic(err)
125	}
126
127	if err := store.Commit(ctx, txn); err != nil {
128		panic(err)
129	}
130
131	var buf bytes.Buffer
132	repl := newRepl(store, &buf)
133	repl.OneShot(ctx, "s = 4")
134	buf.Reset()
135
136	result := repl.complete("")
137	expected := []string{
138		"data.a.b.c.p",
139		"data.a.b.c.q",
140		"data.a.b.d.r",
141		"data.repl.s",
142		"data.repl.version",
143	}
144
145	sort.Strings(result)
146	sort.Strings(expected)
147
148	if !reflect.DeepEqual(result, expected) {
149		t.Fatalf("Expected %v but got: %v", expected, result)
150	}
151
152	result = repl.complete("data.a.b")
153	expected = []string{
154		"data.a.b.c.p",
155		"data.a.b.c.q",
156		"data.a.b.d.r",
157	}
158
159	sort.Strings(result)
160	sort.Strings(expected)
161
162	if !reflect.DeepEqual(result, expected) {
163		t.Fatalf("Expected %v but got: %v", expected, result)
164	}
165
166	result = repl.complete("data.a.b.c.p[x]")
167	expected = []string{}
168
169	if !reflect.DeepEqual(result, expected) {
170		t.Fatalf("Expected %v but got: %v", expected, result)
171	}
172
173	repl.OneShot(ctx, "import data.a.b.c.p as xyz")
174	repl.OneShot(ctx, "import data.a.b.d")
175
176	result = repl.complete("x")
177	expected = []string{
178		"xyz",
179	}
180
181	if !reflect.DeepEqual(result, expected) {
182		t.Fatalf("Expected %v but got: %v", expected, result)
183	}
184}
185
186func TestDump(t *testing.T) {
187	ctx := context.Background()
188	input := `{"a": [1,2,3,4]}`
189	var data map[string]interface{}
190	err := util.UnmarshalJSON([]byte(input), &data)
191	if err != nil {
192		panic(err)
193	}
194	store := inmem.NewFromObject(data)
195	var buffer bytes.Buffer
196	repl := newRepl(store, &buffer)
197	repl.OneShot(ctx, "dump")
198	expectOutput(t, buffer.String(), "{\"a\":[1,2,3,4]}\n")
199}
200
201func TestDumpPath(t *testing.T) {
202	ctx := context.Background()
203	input := `{"a": [1,2,3,4]}`
204	var data map[string]interface{}
205	err := util.UnmarshalJSON([]byte(input), &data)
206	if err != nil {
207		panic(err)
208	}
209	store := inmem.NewFromObject(data)
210	var buffer bytes.Buffer
211	repl := newRepl(store, &buffer)
212
213	dir, err := ioutil.TempDir("", "dump-path-test")
214	if err != nil {
215		t.Fatal(err)
216	}
217
218	defer os.RemoveAll(dir)
219	file := filepath.Join(dir, "tmpfile")
220	repl.OneShot(ctx, fmt.Sprintf("dump %s", file))
221
222	if buffer.String() != "" {
223		t.Errorf("Expected no output but got: %v", buffer.String())
224	}
225
226	bs, err := ioutil.ReadFile(file)
227	if err != nil {
228		t.Fatalf("Expected file read to succeed but got: %v", err)
229	}
230
231	var result map[string]interface{}
232	if err := util.UnmarshalJSON(bs, &result); err != nil {
233		t.Fatalf("Expected json unmarshal to succeed but got: %v", err)
234	}
235
236	if !reflect.DeepEqual(data, result) {
237		t.Fatalf("Expected dumped json to equal %v but got: %v", data, result)
238	}
239}
240
241func TestHelp(t *testing.T) {
242	topics["deadbeef"] = topicDesc{
243		fn: func(w io.Writer) error {
244			fmt.Fprintln(w, "blah blah blah")
245			return nil
246		},
247	}
248
249	ctx := context.Background()
250	store := inmem.New()
251	var buffer bytes.Buffer
252	repl := newRepl(store, &buffer)
253	repl.OneShot(ctx, "help deadbeef")
254
255	expected := "blah blah blah\n"
256
257	if buffer.String() != expected {
258		t.Fatalf("Unexpected output from help topic: %v", buffer.String())
259	}
260}
261
262func TestShow(t *testing.T) {
263	ctx := context.Background()
264	store := inmem.New()
265	var buffer bytes.Buffer
266	repl := newRepl(store, &buffer)
267
268	repl.OneShot(ctx, `package repl_test`)
269	repl.OneShot(ctx, "show")
270	assertREPLText(t, buffer, "package repl_test\n\n")
271	buffer.Reset()
272
273	repl.OneShot(ctx, "import input.xyz")
274	repl.OneShot(ctx, "show")
275
276	expected := `package repl_test
277
278import input.xyz` + "\n\n"
279	assertREPLText(t, buffer, expected)
280	buffer.Reset()
281
282	repl.OneShot(ctx, "import data.foo as bar")
283	repl.OneShot(ctx, "show")
284
285	expected = `package repl_test
286
287import data.foo as bar
288import input.xyz` + "\n\n"
289	assertREPLText(t, buffer, expected)
290	buffer.Reset()
291
292	repl.OneShot(ctx, `p[1] { true }`)
293	repl.OneShot(ctx, `p[2] { true }`)
294	repl.OneShot(ctx, "show")
295
296	expected = `package repl_test
297
298import data.foo as bar
299import input.xyz
300
301p[1]
302
303p[2]` + "\n\n"
304	assertREPLText(t, buffer, expected)
305	buffer.Reset()
306
307	repl.OneShot(ctx, "package abc")
308	repl.OneShot(ctx, "show")
309
310	assertREPLText(t, buffer, "package abc\n\n")
311	buffer.Reset()
312
313	repl.OneShot(ctx, "package repl_test")
314	repl.OneShot(ctx, "show")
315
316	assertREPLText(t, buffer, expected)
317	buffer.Reset()
318}
319
320func TestTypes(t *testing.T) {
321	ctx := context.Background()
322	store := inmem.New()
323	var buffer bytes.Buffer
324	repl := newRepl(store, &buffer)
325
326	repl.OneShot(ctx, "types")
327	repl.OneShot(ctx, "data.repl.version[x]")
328
329	output := strings.TrimSpace(buffer.String())
330
331	exp := []string{
332		"# data.repl.version[x]: string",
333		"# x: string",
334	}
335
336	for i := range exp {
337		if !strings.Contains(output, exp[i]) {
338			t.Fatalf("Expected output to contain %q but got: %v", exp[i], output)
339		}
340	}
341
342}
343
344func TestUnknown(t *testing.T) {
345	ctx := context.Background()
346	store := inmem.New()
347	var buffer bytes.Buffer
348	repl := newRepl(store, &buffer)
349
350	repl.OneShot(ctx, "xs = [1,2,3]")
351
352	err := repl.OneShot(ctx, "unknown input")
353	if err != nil {
354		t.Fatal("Unexpected command error:", err)
355	}
356
357	repl.OneShot(ctx, "data.repl.xs[i] = x; input.x = x")
358
359	output := strings.TrimSpace(buffer.String())
360	expected := strings.TrimSpace(`
361input.x = 1; i = 0; x = 1
362input.x = 2; i = 1; x = 2
363input.x = 3; i = 2; x = 3
364`)
365
366	if output != expected {
367		t.Fatalf("Unexpected output. Expected:\n\n%v\n\nGot:\n\n%v", expected, output)
368	}
369}
370
371func TestUnset(t *testing.T) {
372	ctx := context.Background()
373	store := inmem.New()
374	var buffer bytes.Buffer
375	repl := newRepl(store, &buffer)
376
377	var err error
378
379	repl.OneShot(ctx, "magic = 23")
380	repl.OneShot(ctx, "p = 3.14")
381	repl.OneShot(ctx, "unset p")
382
383	err = repl.OneShot(ctx, "p")
384
385	if _, ok := err.(ast.Errors); !ok {
386		t.Fatalf("Expected AST error but got: %v", err)
387	}
388
389	buffer.Reset()
390	repl.OneShot(ctx, "p = 3.14")
391	repl.OneShot(ctx, `p = 3 { false }`)
392	repl.OneShot(ctx, "unset p")
393
394	err = repl.OneShot(ctx, "p")
395	if _, ok := err.(ast.Errors); !ok {
396		t.Fatalf("Expected AST error but got err: %v, output: %v", err, buffer.String())
397	}
398
399	if err := repl.OneShot(ctx, "unset "); err == nil {
400		t.Fatalf("Expected unset error for bad syntax but got: %v", buffer.String())
401	}
402
403	if err := repl.OneShot(ctx, "unset 1=1"); err == nil {
404		t.Fatalf("Expected unset error for bad syntax but got: %v", buffer.String())
405	}
406
407	if err := repl.OneShot(ctx, `unset "p"`); err == nil {
408		t.Fatalf("Expected unset error for bad syntax but got: %v", buffer.String())
409	}
410
411	buffer.Reset()
412	repl.OneShot(ctx, "p(x) = y { y = x }")
413	repl.OneShot(ctx, "unset p")
414
415	err = repl.OneShot(ctx, "data.repl.p(1, 2)")
416	if err == nil || err.Error() != `1 error occurred: 1:1: rego_type_error: undefined function data.repl.p` {
417		t.Fatalf("Expected eval error (undefined built-in) but got err: '%v'", err)
418	}
419
420	buffer.Reset()
421	repl.OneShot(ctx, "p(1, x) = y { y = x }")
422	repl.OneShot(ctx, "p(2, x) = y { y = x+1 }")
423	repl.OneShot(ctx, "unset p")
424
425	err = repl.OneShot(ctx, "data.repl.p(1, 2, 3)")
426	if err == nil || err.Error() != `1 error occurred: 1:1: rego_type_error: undefined function data.repl.p` {
427		t.Fatalf("Expected eval error (undefined built-in) but got err: '%v'", err)
428	}
429
430	buffer.Reset()
431	repl.OneShot(ctx, `unset q`)
432	if buffer.String() != "warning: no matching rules in current module\n" {
433		t.Fatalf("Expected unset error for missing rule but got: %v", buffer.String())
434	}
435
436	buffer.Reset()
437	repl.OneShot(ctx, `unset q`)
438	if buffer.String() != "warning: no matching rules in current module\n" {
439		t.Fatalf("Expected unset error for missing function but got: %v", buffer.String())
440	}
441
442	buffer.Reset()
443	repl.OneShot(ctx, `magic`)
444	if buffer.String() != "23\n" {
445		t.Fatalf("Expected magic to be defined but got: %v", buffer.String())
446	}
447
448	buffer.Reset()
449	repl.OneShot(ctx, `package data.other`)
450	err = repl.OneShot(ctx, `unset magic`)
451	if buffer.String() != "warning: no matching rules in current module\n" {
452		t.Fatalf("Expected unset error for bad syntax but got: %v", buffer.String())
453	}
454
455	repl.OneShot(ctx, `input = {}`)
456
457	if err := repl.OneShot(ctx, `unset input`); err != nil {
458		t.Fatalf("Expected unset to succeed for input: %v", err)
459	}
460
461	buffer.Reset()
462	repl.OneShot(ctx, `not input`)
463
464	if buffer.String() != "true\n" {
465		t.Fatalf("Expected unset input to remove input document: %v", buffer.String())
466	}
467
468}
469
470func TestOneShotEmptyBufferOneExpr(t *testing.T) {
471	ctx := context.Background()
472	store := newTestStore()
473	var buffer bytes.Buffer
474	repl := newRepl(store, &buffer)
475	repl.OneShot(ctx, "data.a[i].b.c[j] = 2")
476	expectOutput(t, buffer.String(), "+---+---+\n| i | j |\n+---+---+\n| 0 | 1 |\n+---+---+\n")
477	buffer.Reset()
478	repl.OneShot(ctx, "data.a[i].b.c[j] = \"deadbeef\"")
479	expectOutput(t, buffer.String(), "undefined\n")
480}
481
482func TestOneShotEmptyBufferOneRule(t *testing.T) {
483	ctx := context.Background()
484	store := newTestStore()
485	var buffer bytes.Buffer
486	repl := newRepl(store, &buffer)
487	repl.OneShot(ctx, `p[x] { data.a[i] = x }`)
488	expectOutput(t, buffer.String(), "")
489}
490
491func TestOneShotBufferedExpr(t *testing.T) {
492	ctx := context.Background()
493	store := newTestStore()
494	var buffer bytes.Buffer
495	repl := newRepl(store, &buffer)
496	repl.OneShot(ctx, "data.a[i].b.c[j] = ")
497	expectOutput(t, buffer.String(), "")
498	repl.OneShot(ctx, "2")
499	expectOutput(t, buffer.String(), "")
500	repl.OneShot(ctx, "")
501	expectOutput(t, buffer.String(), "+---+---+\n| i | j |\n+---+---+\n| 0 | 1 |\n+---+---+\n")
502}
503
504func TestOneShotBufferedRule(t *testing.T) {
505	ctx := context.Background()
506	store := newTestStore()
507	var buffer bytes.Buffer
508	repl := newRepl(store, &buffer)
509	repl.OneShot(ctx, "p[x] { ")
510	expectOutput(t, buffer.String(), "")
511	repl.OneShot(ctx, "data.a[i].b.c[1]")
512	expectOutput(t, buffer.String(), "")
513	repl.OneShot(ctx, " = ")
514	expectOutput(t, buffer.String(), "")
515	repl.OneShot(ctx, "x")
516	expectOutput(t, buffer.String(), "")
517	repl.OneShot(ctx, "}")
518	expectOutput(t, buffer.String(), "")
519	repl.OneShot(ctx, "")
520	expectOutput(t, buffer.String(), "")
521	repl.OneShot(ctx, "p[2]")
522	expectOutput(t, buffer.String(), "2\n")
523}
524
525func TestOneShotJSON(t *testing.T) {
526	ctx := context.Background()
527	store := newTestStore()
528	var buffer bytes.Buffer
529	repl := newRepl(store, &buffer)
530	repl.outputFormat = "json"
531	repl.OneShot(ctx, "data.a[i] = x")
532	var expected interface{}
533	input := `
534	[
535		{
536			"i": 0,
537			"x": {
538			"b": {
539				"c": [
540				true,
541				2,
542				false
543				]
544			}
545			}
546		},
547		{
548			"i": 1,
549			"x": {
550			"b": {
551				"c": [
552				false,
553				true,
554				1
555				]
556			}
557			}
558		}
559	]
560	`
561	if err := util.UnmarshalJSON([]byte(input), &expected); err != nil {
562		panic(err)
563	}
564
565	var result interface{}
566
567	if err := util.UnmarshalJSON(buffer.Bytes(), &result); err != nil {
568		t.Errorf("Unexpected output format: %v", err)
569		return
570	}
571	if !reflect.DeepEqual(expected, result) {
572		t.Errorf("Expected %v but got: %v", expected, result)
573	}
574}
575
576func TestEvalData(t *testing.T) {
577	ctx := context.Background()
578	store := newTestStore()
579	var buffer bytes.Buffer
580	repl := newRepl(store, &buffer)
581	testMod := []byte(`package ex
582
583p = [1, 2, 3] { true }`)
584
585	txn := storage.NewTransactionOrDie(ctx, store, storage.WriteParams)
586
587	if err := store.UpsertPolicy(ctx, txn, "test", testMod); err != nil {
588		panic(err)
589	}
590
591	if err := store.Commit(ctx, txn); err != nil {
592		panic(err)
593	}
594
595	repl.OneShot(ctx, "data")
596
597	expected := parseJSON(`
598	{
599		"a": [
600			{
601			"b": {
602				"c": [
603				true,
604				2,
605				false
606				]
607			}
608			},
609			{
610			"b": {
611				"c": [
612				false,
613				true,
614				1
615				]
616			}
617			}
618		],
619		"ex": {
620			"p": [
621			1,
622			2,
623			3
624			]
625		}
626	}`)
627	result := parseJSON(buffer.String())
628
629	// Strip REPL documents out as these change depending on build settings.
630	data := result.(map[string]interface{})
631	delete(data, "repl")
632
633	if !reflect.DeepEqual(result, expected) {
634		t.Fatalf("Expected:\n%v\n\nGot:\n%v", expected, result)
635	}
636}
637
638func TestEvalFalse(t *testing.T) {
639	ctx := context.Background()
640	store := newTestStore()
641	var buffer bytes.Buffer
642	repl := newRepl(store, &buffer)
643	repl.OneShot(ctx, "false")
644	result := buffer.String()
645	if result != "false\n" {
646		t.Errorf("Expected result to be false but got: %v", result)
647	}
648}
649
650func TestEvalConstantRule(t *testing.T) {
651	ctx := context.Background()
652	store := newTestStore()
653	var buffer bytes.Buffer
654	repl := newRepl(store, &buffer)
655	repl.OneShot(ctx, "pi = 3.14")
656	result := buffer.String()
657	if result != "" {
658		t.Errorf("Expected rule to be defined but got: %v", result)
659		return
660	}
661	buffer.Reset()
662	repl.OneShot(ctx, "pi")
663	result = buffer.String()
664	expected := "3.14\n"
665	if result != expected {
666		t.Errorf("Expected pi to evaluate to 3.14 but got: %v", result)
667		return
668	}
669	buffer.Reset()
670	err := repl.OneShot(ctx, "pi.deadbeef")
671	result = buffer.String()
672	if result != "" || !strings.Contains(err.Error(), "undefined ref: data.repl.pi.deadbeef") {
673		t.Fatalf("Expected pi.deadbeef to fail/error but got:\nresult: %q\nerr: %v", result, err)
674	}
675	buffer.Reset()
676	repl.OneShot(ctx, "pi > 3")
677	result = buffer.String()
678	if result != "true\n" {
679		t.Errorf("Expected pi > 3 to be true but got: %v", result)
680		return
681	}
682}
683
684func TestEvalConstantRuleAssignment(t *testing.T) {
685	ctx := context.Background()
686	store := newTestStore()
687	var buffer bytes.Buffer
688
689	repl := newRepl(store, &buffer)
690	repl.OneShot(ctx, "x = 1")
691	repl.OneShot(ctx, "x := 2")
692	repl.OneShot(ctx, "x := 3")
693	repl.OneShot(ctx, "x")
694	result := buffer.String()
695	if result != "3\n" {
696		t.Fatalf("Expected 3 but got: %v", result)
697	}
698
699	buffer.Reset()
700	repl.OneShot(ctx, "x = 3")
701	result = buffer.String()
702	if result != "true\n" {
703		t.Fatalf("Expected true but got: %v", result)
704	}
705
706	buffer.Reset()
707	repl.OneShot(ctx, "input = 0")
708	repl.OneShot(ctx, "input := 1")
709	repl.OneShot(ctx, "input")
710	result = buffer.String()
711	if result != "1\n" {
712		t.Fatalf("Expected 1 but got: %v", result)
713	}
714}
715
716func TestEvalSingleTermMultiValue(t *testing.T) {
717	ctx := context.Background()
718	store := newTestStore()
719	var buffer bytes.Buffer
720	repl := newRepl(store, &buffer)
721	repl.outputFormat = "json"
722
723	input := `
724	[
725		{
726			"i": 0
727		},
728		{
729			"i": 0
730		},
731		{
732			"i": 1
733		},
734		{
735			"i": 1
736		}
737	]`
738
739	var expected interface{}
740	if err := util.UnmarshalJSON([]byte(input), &expected); err != nil {
741		panic(err)
742	}
743
744	repl.OneShot(ctx, "data.a[i].b.c[_]")
745	var result interface{}
746	if err := util.UnmarshalJSON(buffer.Bytes(), &result); err != nil {
747		t.Errorf("Expected valid JSON document: %v: %v", err, buffer.String())
748		return
749	}
750
751	if !reflect.DeepEqual(expected, result) {
752		t.Errorf("Expected %v but got: %v", expected, result)
753		return
754	}
755
756	buffer.Reset()
757
758	repl.OneShot(ctx, "data.deadbeef[x]")
759	s := buffer.String()
760	if s != "undefined\n" {
761		t.Errorf("Expected undefined from reference but got: %v", s)
762		return
763	}
764
765	buffer.Reset()
766
767	repl.OneShot(ctx, `p[x] { a = [1, 2, 3, 4]; a[_] = x }`)
768	buffer.Reset()
769	repl.OneShot(ctx, "p[x]")
770
771	input = `
772	[
773		{
774			"x": 1
775		},
776		{
777			"x": 2
778		},
779		{
780			"x": 3
781		},
782		{
783			"x": 4
784		}
785	]
786	`
787
788	if err := util.UnmarshalJSON([]byte(input), &expected); err != nil {
789		panic(err)
790	}
791
792	if err := util.UnmarshalJSON(buffer.Bytes(), &result); err != nil {
793		t.Errorf("Expected valid JSON document: %v: %v", err, buffer.String())
794		return
795	}
796
797	if !reflect.DeepEqual(expected, result) {
798		t.Errorf("Exepcted %v but got: %v", expected, result)
799	}
800}
801
802func TestEvalSingleTermMultiValueSetRef(t *testing.T) {
803	ctx := context.Background()
804	store := newTestStore()
805	var buffer bytes.Buffer
806	repl := newRepl(store, &buffer)
807	repl.outputFormat = "json"
808	repl.OneShot(ctx, `p[1] { true }`)
809	repl.OneShot(ctx, `p[2] { true }`)
810	repl.OneShot(ctx, `q = {3, 4} { true }`)
811	repl.OneShot(ctx, `r = [x, y] { x = {5, 6}; y = [7, 8] }`)
812
813	repl.OneShot(ctx, "p[x]")
814	expected := parseJSON(`[{"x": 1}, {"x": 2}]`)
815	result := parseJSON(buffer.String())
816	if !reflect.DeepEqual(result, expected) {
817		t.Fatalf("Expected %v but got: %v", expected, result)
818	}
819
820	buffer.Reset()
821	repl.OneShot(ctx, "q[x]")
822	expected = parseJSON(`[{"x": 3}, {"x": 4}]`)
823	result = parseJSON(buffer.String())
824	if !reflect.DeepEqual(result, expected) {
825		t.Fatalf("Expected %v but got: %v", expected, result)
826	}
827
828	// Example below shows behavior for ref that iterates an embedded set. The
829	// tricky part here is that r[_] may refer to multiple collection types. If
830	// we eventually have a way of distinguishing between the bindings added for
831	// refs to sets, then those bindings could be filtered out. For now this is
832	// acceptable, as it should be an edge case.
833	buffer.Reset()
834	repl.OneShot(ctx, "r[_][x]")
835	expected = parseJSON(`[{"x": 5}, {"x": 6}, {"x": 0}, {"x": 1}]`)
836	result = parseJSON(buffer.String())
837	if !reflect.DeepEqual(result, expected) {
838		t.Fatalf("Expected %v but got: %v", expected, result)
839	}
840}
841
842func TestEvalRuleCompileError(t *testing.T) {
843	ctx := context.Background()
844	store := newTestStore()
845	var buffer bytes.Buffer
846	repl := newRepl(store, &buffer)
847	err := repl.OneShot(ctx, `p[x] { true }`)
848	expected := "x is unsafe"
849	if !strings.Contains(err.Error(), expected) {
850		t.Errorf("Expected error to contain %v but got: %v (err: %v)", expected, buffer.String(), err)
851		return
852	}
853	buffer.Reset()
854	err = repl.OneShot(ctx, `p = true { true }`)
855	result := buffer.String()
856	if err != nil || result != "" {
857		t.Errorf("Expected valid rule to compile (because state should be unaffected) but got: %v (err: %v)", result, err)
858	}
859}
860
861func TestEvalBodyCompileError(t *testing.T) {
862	ctx := context.Background()
863	store := newTestStore()
864	var buffer bytes.Buffer
865	repl := newRepl(store, &buffer)
866	repl.outputFormat = "json"
867	err := repl.OneShot(ctx, `x = 1; y > x`)
868	if _, ok := err.(ast.Errors); !ok {
869		t.Fatalf("Expected error message in output but got`: %v", buffer.String())
870	}
871	buffer.Reset()
872	repl.OneShot(ctx, `x = 1; y = 2; y > x`)
873	var result2 []interface{}
874	err = util.UnmarshalJSON(buffer.Bytes(), &result2)
875	if err != nil {
876		t.Errorf("Expected valid JSON output but got: %v", buffer.String())
877		return
878	}
879	expected2 := []interface{}{
880		map[string]interface{}{
881			"x": json.Number("1"),
882			"y": json.Number("2"),
883		},
884	}
885	if !reflect.DeepEqual(expected2, result2) {
886		t.Errorf(`Expected [{"x": 1, "y": 2}] but got: %v"`, result2)
887		return
888	}
889}
890
891func TestEvalBodyContainingWildCards(t *testing.T) {
892	ctx := context.Background()
893	store := newTestStore()
894	var buffer bytes.Buffer
895	repl := newRepl(store, &buffer)
896	repl.OneShot(ctx, "data.a[_].b.c[_] = x")
897	expected := strings.TrimSpace(`
898+-------+
899|   x   |
900+-------+
901| true  |
902| 2     |
903| false |
904| false |
905| true  |
906| 1     |
907+-------+`)
908	result := strings.TrimSpace(buffer.String())
909	if result != expected {
910		t.Errorf("Expected only a single column of output but got:\n%v", result)
911	}
912
913}
914
915func TestEvalBodyInput(t *testing.T) {
916	ctx := context.Background()
917	store := newTestStore()
918	var buffer bytes.Buffer
919	repl := newRepl(store, &buffer)
920
921	repl.OneShot(ctx, `package repl`)
922	repl.OneShot(ctx, `input["foo.bar"] = "hello" { true }`)
923	repl.OneShot(ctx, `input["baz"] = data.a[0].b.c[2] { true }`)
924	repl.OneShot(ctx, `package test`)
925	repl.OneShot(ctx, "import input.baz")
926	repl.OneShot(ctx, `p = true { input["foo.bar"] = "hello"; baz = false }`)
927	repl.OneShot(ctx, "p")
928
929	result := buffer.String()
930	if result != "true\n" {
931		t.Fatalf("expected true but got: %v", result)
932	}
933}
934
935func TestEvalBodyInputComplete(t *testing.T) {
936	ctx := context.Background()
937	store := newTestStore()
938	var buffer bytes.Buffer
939	repl := newRepl(store, &buffer)
940
941	// Test that input can be defined completely:
942	// https://github.com/open-policy-agent/opa/issues/231
943	repl.OneShot(ctx, `package repl`)
944	repl.OneShot(ctx, `input = 1`)
945	repl.OneShot(ctx, `input`)
946
947	result := buffer.String()
948	if result != "1\n" {
949		t.Fatalf("Expected 1 but got: %v", result)
950	}
951
952	buffer.Reset()
953
954	// Test that input is as expected
955	repl.OneShot(ctx, `package ex1`)
956	repl.OneShot(ctx, `x = input`)
957	repl.OneShot(ctx, `x`)
958
959	result = buffer.String()
960	if result != "1\n" {
961		t.Fatalf("Expected 1 but got: %v", result)
962	}
963
964	buffer.Reset()
965
966	// Test that local input replaces other inputs
967	repl.OneShot(ctx, `package ex2`)
968	repl.OneShot(ctx, `input = 2`)
969	repl.OneShot(ctx, `input`)
970
971	result = buffer.String()
972
973	if result != "2\n" {
974		t.Fatalf("Expected 2 but got: %v", result)
975	}
976
977	buffer.Reset()
978
979	// Test that original input is intact
980	repl.OneShot(ctx, `package ex3`)
981	repl.OneShot(ctx, `input`)
982
983	result = buffer.String()
984
985	if result != "1\n" {
986		t.Fatalf("Expected 1 but got: %v", result)
987	}
988
989	// Test that deferencing undefined input results in undefined
990	buffer.Reset()
991
992	repl = newRepl(store, &buffer)
993	repl.OneShot(ctx, `input.p`)
994	result = buffer.String()
995	if result != "undefined\n" {
996		t.Fatalf("Expected undefined but got: %v", result)
997	}
998
999	buffer.Reset()
1000	repl.OneShot(ctx, `input.p = false`)
1001	result = buffer.String()
1002	if result != "undefined\n" {
1003		t.Fatalf("Expected undefined but got: %v", result)
1004	}
1005
1006}
1007
1008func TestEvalBodyWith(t *testing.T) {
1009	ctx := context.Background()
1010	store := newTestStore()
1011	var buffer bytes.Buffer
1012	repl := newRepl(store, &buffer)
1013
1014	repl.OneShot(ctx, `p = true { input.foo = "bar" }`)
1015	repl.OneShot(ctx, "p")
1016
1017	if buffer.String() != "undefined\n" {
1018		t.Fatalf("Expected undefined but got: %v", buffer.String())
1019	}
1020
1021	buffer.Reset()
1022
1023	repl.OneShot(ctx, `p with input.foo as "bar"`)
1024
1025	result := buffer.String()
1026	expected := "true\n"
1027
1028	if result != expected {
1029		t.Fatalf("Expected true but got: %v", result)
1030	}
1031}
1032
1033func TestEvalBodyRewrittenBuiltin(t *testing.T) {
1034	ctx := context.Background()
1035	store := newTestStore()
1036	var buffer bytes.Buffer
1037	repl := newRepl(store, &buffer)
1038	repl.OneShot(ctx, "json")
1039	repl.OneShot(ctx, `p[x] { a[x]; a = [1,2,3,4] }`)
1040	repl.OneShot(ctx, "p[x] > 1")
1041	result := util.MustUnmarshalJSON(buffer.Bytes())
1042	expected := util.MustUnmarshalJSON([]byte(`[{"x": 2}, {"x": 3}]`))
1043	if util.Compare(result, expected) != 0 {
1044		t.Fatalf("Expected %v but got: %v", expected, result)
1045	}
1046}
1047
1048func TestEvalBodyRewrittenRef(t *testing.T) {
1049	ctx := context.Background()
1050	store := newTestStore()
1051	var buffer bytes.Buffer
1052	repl := newRepl(store, &buffer)
1053	repl.OneShot(ctx, "json")
1054	repl.OneShot(ctx, `i = 1`)
1055	repl.OneShot(ctx, `data.a[0].b.c[i]`)
1056	result := util.MustUnmarshalJSON(buffer.Bytes())
1057	expected := util.MustUnmarshalJSON([]byte(`2`))
1058	if util.Compare(result, expected) != 0 {
1059		t.Fatalf("Expected %v but got: %v", expected, result)
1060	}
1061	buffer.Reset()
1062	repl.OneShot(ctx, "p = {1,2,3}")
1063	repl.OneShot(ctx, "p")
1064	result = util.MustUnmarshalJSON(buffer.Bytes())
1065	expected = util.MustUnmarshalJSON([]byte(`[1,2,3]`))
1066	if util.Compare(result, expected) != 0 {
1067		t.Fatalf("Expected %v but got: %v", expected, result)
1068	}
1069	buffer.Reset()
1070	repl.OneShot(ctx, "p[x]")
1071	result = util.MustUnmarshalJSON(buffer.Bytes())
1072	expected = util.MustUnmarshalJSON([]byte(`[{"x": 1}, {"x": 2}, {"x": 3}]`))
1073	if util.Compare(result, expected) != 0 {
1074		t.Fatalf("Expected %v but got: %v", expected, result)
1075	}
1076}
1077
1078func TestEvalImport(t *testing.T) {
1079	ctx := context.Background()
1080	store := newTestStore()
1081	var buffer bytes.Buffer
1082	repl := newRepl(store, &buffer)
1083	repl.OneShot(ctx, "import data.a")
1084	if len(buffer.Bytes()) != 0 {
1085		t.Errorf("Expected no output but got: %v", buffer.String())
1086		return
1087	}
1088	buffer.Reset()
1089	repl.OneShot(ctx, "a[0].b.c[0] = true")
1090	result := buffer.String()
1091	expected := "true\n"
1092	if result != expected {
1093		t.Errorf("Expected expression to evaluate successfully but got: %v", result)
1094		return
1095	}
1096
1097	// https://github.com/open-policy-agent/opa/issues/158 - re-run query to
1098	// make sure import is not lost
1099	buffer.Reset()
1100	repl.OneShot(ctx, "a[0].b.c[0] = true")
1101	result = buffer.String()
1102	expected = "true\n"
1103	if result != expected {
1104		t.Fatalf("Expected expression to evaluate successfully but got: %v", result)
1105	}
1106}
1107
1108func TestEvalPackage(t *testing.T) {
1109	ctx := context.Background()
1110	store := newTestStore()
1111	var buffer bytes.Buffer
1112	repl := newRepl(store, &buffer)
1113	repl.OneShot(ctx, `package foo.bar`)
1114	repl.OneShot(ctx, `p = true { true }`)
1115	repl.OneShot(ctx, `package baz.qux`)
1116	buffer.Reset()
1117	err := repl.OneShot(ctx, "p")
1118	if !strings.Contains(err.Error(), "p is unsafe") {
1119		t.Fatalf("Expected unsafe variable error but got: %v", err)
1120	}
1121	repl.OneShot(ctx, "import data.foo.bar.p")
1122	buffer.Reset()
1123	repl.OneShot(ctx, "p")
1124	if buffer.String() != "true\n" {
1125		t.Errorf("Expected expression to eval successfully but got: %v", buffer.String())
1126		return
1127	}
1128}
1129
1130func TestMetrics(t *testing.T) {
1131	ctx := context.Background()
1132	store := newTestStore()
1133	var buffer bytes.Buffer
1134
1135	repl := newRepl(store, &buffer)
1136	repl.OneShot(ctx, "a = {[1,2], [3,4]}")
1137	repl.OneShot(ctx, "metrics")
1138	repl.OneShot(ctx, `[x | a[x]]`)
1139	if !strings.Contains(buffer.String(), "timer_rego_query_compile_ns") {
1140		t.Fatal("Expected output to contain well known metric key but got:", buffer.String())
1141	}
1142
1143	buffer.Reset()
1144	repl.OneShot(ctx, `[x | a[x]]`)
1145	if !strings.Contains(buffer.String(), "timer_rego_query_compile_ns") {
1146		t.Fatal("Expected output to contain well known metric key but got:", buffer.String())
1147	}
1148
1149	buffer.Reset()
1150	repl.OneShot(ctx, "metrics")
1151	repl.OneShot(ctx, `[x | a[x]]`)
1152
1153	expected := `[
1154  [
1155    1,
1156    2
1157  ],
1158  [
1159    3,
1160    4
1161  ]
1162]
1163`
1164
1165	if expected != buffer.String() {
1166		t.Fatalf("Expected output to be exactly:\n%v\n\nGot:\n\n%v\n", expected, buffer.String())
1167	}
1168}
1169
1170func TestInstrument(t *testing.T) {
1171	ctx := context.Background()
1172	store := newTestStore()
1173	var buffer bytes.Buffer
1174
1175	repl := newRepl(store, &buffer)
1176
1177	// Turn on instrumentation w/o turning on metrics.
1178	repl.OneShot(ctx, "instrument")
1179	repl.OneShot(ctx, "true")
1180
1181	result := buffer.String()
1182
1183	if !strings.Contains(result, "histogram_eval_op_plug") {
1184		t.Fatal("Expected plug histogram in output but got:", result)
1185	}
1186
1187	buffer.Reset()
1188
1189	// Turn off instrumentation.
1190	repl.OneShot(ctx, "instrument")
1191	repl.OneShot(ctx, "true")
1192
1193	result = buffer.String()
1194
1195	if strings.Contains(result, "histogram_eval_op_plug") {
1196		t.Fatal("Expected instrumentation to be turned off but got:", result)
1197	}
1198
1199	buffer.Reset()
1200
1201	// Turn on metrics and then turn on instrumentation.
1202	repl.OneShot(ctx, "metrics")
1203	repl.OneShot(ctx, "true")
1204
1205	result = buffer.String()
1206
1207	if strings.Contains(result, "histogram_eval_op_plug") {
1208		t.Fatal("Expected instrumentation to be turned off but got:", result)
1209	}
1210
1211	if !strings.Contains(result, "timer_rego_query_eval_ns") {
1212		t.Fatal("Expected metrics to be turned on but got:", result)
1213	}
1214
1215	buffer.Reset()
1216
1217	repl.OneShot(ctx, "instrument")
1218	repl.OneShot(ctx, "true")
1219
1220	result = buffer.String()
1221
1222	if !strings.Contains(result, "histogram_eval_op_plug") {
1223		t.Fatal("Expected instrumentation to be turned on but got:", result)
1224	}
1225
1226	if !strings.Contains(result, "timer_rego_query_eval_ns") {
1227		t.Fatal("Expected metrics to be turned on but got:", result)
1228	}
1229
1230}
1231
1232func TestEvalTrace(t *testing.T) {
1233	ctx := context.Background()
1234	store := newTestStore()
1235	var buffer bytes.Buffer
1236	repl := newRepl(store, &buffer)
1237	repl.OneShot(ctx, "trace")
1238	repl.OneShot(ctx, `data.a[i].b.c[j] = x; data.a[k].b.c[x] = 1`)
1239	expected := strings.TrimSpace(`
1240Enter data.a[i].b.c[j] = x; data.a[k].b.c[x] = 1
1241| Eval data.a[i].b.c[j] = x
1242| Eval data.a[k].b.c[x] = 1
1243| Fail data.a[k].b.c[x] = 1
1244| Redo data.a[i].b.c[j] = x
1245| Eval data.a[k].b.c[x] = 1
1246| Exit data.a[i].b.c[j] = x; data.a[k].b.c[x] = 1
1247Redo data.a[i].b.c[j] = x; data.a[k].b.c[x] = 1
1248| Redo data.a[k].b.c[x] = 1
1249| Redo data.a[i].b.c[j] = x
1250| Eval data.a[k].b.c[x] = 1
1251| Fail data.a[k].b.c[x] = 1
1252| Redo data.a[i].b.c[j] = x
1253| Eval data.a[k].b.c[x] = 1
1254| Fail data.a[k].b.c[x] = 1
1255| Redo data.a[i].b.c[j] = x
1256| Eval data.a[k].b.c[x] = 1
1257| Fail data.a[k].b.c[x] = 1
1258| Redo data.a[i].b.c[j] = x
1259| Eval data.a[k].b.c[x] = 1
1260| Fail data.a[k].b.c[x] = 1
1261| Redo data.a[i].b.c[j] = x
1262+---+---+---+---+
1263| i | j | k | x |
1264+---+---+---+---+
1265| 0 | 1 | 1 | 2 |
1266+---+---+---+---+`)
1267	expected += "\n"
1268
1269	if expected != buffer.String() {
1270		t.Fatalf("Expected output to be exactly:\n%v\n\nGot:\n\n%v\n", expected, buffer.String())
1271	}
1272}
1273
1274func TestTruncatePrettyOutput(t *testing.T) {
1275	ctx := context.Background()
1276	store := inmem.New()
1277	var buffer bytes.Buffer
1278	repl := newRepl(store, &buffer)
1279	repl.prettyLimit = 1000 // crank up limit to test repl command
1280	repl.OneShot(ctx, "pretty-limit 80")
1281	repl.OneShot(ctx, "data[x]")
1282	for _, line := range strings.Split(buffer.String(), "\n") {
1283		// | "repl" | {"version": <elided>... |
1284		if len(line) > 96 {
1285			t.Fatalf("Expected len(line) to be < 96 but got:\n\n%v", buffer)
1286		}
1287	}
1288	buffer.Reset()
1289	if err := repl.OneShot(ctx, "pretty-limit"); err == nil || !strings.Contains(err.Error(), "usage: pretty-limit <n>") {
1290		t.Fatalf("Expected usage error but got: %v", err)
1291	}
1292}
1293
1294func assertREPLText(t *testing.T, buf bytes.Buffer, expected string) {
1295	result := buf.String()
1296	if result != expected {
1297		t.Fatalf("Expected:\n%v\n\nString:\n\n%v\nGot:\n%v\n\nString:\n\n%v", []byte(expected), expected, []byte(result), result)
1298	}
1299}
1300
1301func expectOutput(t *testing.T, output string, expected string) {
1302	if output != expected {
1303		t.Errorf("Repl output: expected %#v but got %#v", expected, output)
1304	}
1305}
1306
1307func newRepl(store storage.Store, buffer *bytes.Buffer) *REPL {
1308	repl := New(store, "", buffer, "", 0, "")
1309	return repl
1310}
1311
1312func newTestStore() storage.Store {
1313	input := `
1314    {
1315        "a": [
1316            {
1317                "b": {
1318                    "c": [true,2,false]
1319                }
1320            },
1321            {
1322                "b": {
1323                    "c": [false,true,1]
1324                }
1325            }
1326        ]
1327    }
1328    `
1329	var data map[string]interface{}
1330	err := util.UnmarshalJSON([]byte(input), &data)
1331	if err != nil {
1332		panic(err)
1333	}
1334	return inmem.NewFromObject(data)
1335}
1336
1337func parseJSON(s string) interface{} {
1338	var v interface{}
1339	if err := util.UnmarshalJSON([]byte(s), &v); err != nil {
1340		panic(err)
1341	}
1342	return v
1343}
1344