From a079e6f810093fa8ef6d97bb8aa828a9f525aef2 Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Wed, 25 Jan 2017 19:54:09 -0800
Subject: [PATCH] Fix input document definition in REPL

Fixes #231
---
 repl/repl.go      | 21 +++++++++++++-----
 repl/repl_test.go | 56 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/repl/repl.go b/repl/repl.go
index 081dbd9e6f..a79c4de765 100644
--- a/repl/repl.go
+++ b/repl/repl.go
@@ -582,21 +582,32 @@ func (r *REPL) evalStatement(ctx context.Context, stmt interface{}) error {
 		if err != nil {
 			return err
 		}
+
 		body, err := r.compileBody(s, input)
+
 		if err != nil {
-			return err
-		}
-		if rule, err := ast.ParseRuleFromBody(body); err == nil {
-			if err := r.compileRule(rule); err != nil {
+			rule, err2 := ast.ParseRuleFromBody(s)
+			if err2 != nil {
+				// The statement cannot be understood as a rule, so the original
+				// error returned from compiling the query should be given the
+				// caller.
 				return err
 			}
-			return nil
+			return r.compileRule(rule)
 		}
+
+		rule, err3 := ast.ParseRuleFromBody(body)
+		if err3 == nil {
+			return r.compileRule(rule)
+		}
+
 		compiler, err := r.loadCompiler()
 		if err != nil {
 			return err
 		}
+
 		return r.evalBody(ctx, compiler, input, body)
+
 	case *ast.Rule:
 		if err := r.compileRule(s); err != nil {
 			fmt.Fprintln(r.output, "error:", err)
diff --git a/repl/repl_test.go b/repl/repl_test.go
index efaf58abef..f5604c2d39 100644
--- a/repl/repl_test.go
+++ b/repl/repl_test.go
@@ -720,6 +720,62 @@ func TestEvalBodyInput(t *testing.T) {
 	}
 }
 
+func TestEvalBodyInputComplete(t *testing.T) {
+	ctx := context.Background()
+	store := newTestStore()
+	var buffer bytes.Buffer
+	repl := newRepl(store, &buffer)
+
+	// Test that input can be defined completely:
+	// https://github.com/open-policy-agent/opa/issues/231
+	repl.OneShot(ctx, `package repl`)
+	repl.OneShot(ctx, `input = 1`)
+	repl.OneShot(ctx, `input`)
+
+	result := buffer.String()
+	if result != "1\n" {
+		t.Fatalf("Expected 1 but got: %v", result)
+	}
+
+	buffer.Reset()
+
+	// Test that input is as expected
+	repl.OneShot(ctx, `package ex1`)
+	repl.OneShot(ctx, `x = input`)
+	repl.OneShot(ctx, `x`)
+
+	result = buffer.String()
+	if result != "1\n" {
+		t.Fatalf("Expected 1 but got: %v", result)
+	}
+
+	buffer.Reset()
+
+	// Test that local input replaces other inputs
+	repl.OneShot(ctx, `package ex2`)
+	repl.OneShot(ctx, `input = 2`)
+	repl.OneShot(ctx, `input`)
+
+	result = buffer.String()
+
+	if result != "2\n" {
+		t.Fatalf("Expected 2 but got: %v", result)
+	}
+
+	buffer.Reset()
+
+	// Test that original input is intact
+	repl.OneShot(ctx, `package ex3`)
+	repl.OneShot(ctx, `input`)
+
+	result = buffer.String()
+
+	if result != "1\n" {
+		t.Fatalf("Expected 1 bu got: %v", result)
+	}
+
+}
+
 func TestEvalImport(t *testing.T) {
 	ctx := context.Background()
 	store := newTestStore()