1 files changed, 95 insertions, 0 deletions
diff --git a/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml b/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
index e69de29..e447473 100644
--- a/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
+++ b/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
@@ -0,0 +1,95 @@
+open Core
+
+module To_test = struct
+  let extract_locals = Compiler_rewrite.Compiler.extract_local_variables_wrapper
+end
+
+let strings = Alcotest.testable
+                (Format.pp_print_list Format.pp_print_string)
+                (List.equal String.equal)
+
+let invoke ?(stdin="") ?(env=[| |]) cmd =
+  let ic, oc, ec = UnixLabels.open_process_full cmd ~env:env in
+  Out_channel.output_string oc stdin ;
+  Out_channel.close oc;
+  (In_channel.input_all ic,
+   In_channel.input_all ec)
+
+let rego_to_ast policy_code =
+  let (stdout, _) = invoke ("/home/jakob/.local/bin/opa parse -f json /dev/stdin") ~stdin:policy_code ~env:[| |] in
+  Yojson.Basic.from_string stdout |> Compiler_rewrite.Compiler.parse_tree
+
+let extract_query policy =
+  match (rego_to_ast policy) with
+  | { policy = [{ body = body; _ }]; _ } -> body
+  | _ -> failwith "failed to parse"
+
+let test_local_extraction_simple_1 () =
+  let policy_code = {|
+package test
+
+q {
+  x := 2
+}
+  |} in
+  Alcotest.(check strings)
+    "expected identifier list"
+    ["x"]
+    (To_test.extract_locals (extract_query policy_code))
+
+let test_local_extraction_simple_2 () =
+  let policy_code = {|
+package test
+
+q {
+  x := 2
+  y := 3
+  z := 4
+}
+  |} in
+  Alcotest.(check strings)
+    "expected identifier list"
+    ["x"; "y"; "z"]
+    (To_test.extract_locals (extract_query policy_code))
+
+let test_local_dependent_on_previously_defined_local () =
+  let policy_code = {|
+package test
+
+q {
+  x := 2
+  y := x
+}
+  |} in
+  Alcotest.(check strings)
+    "expected identifier list"
+    ["x"; "y"]
+    (To_test.extract_locals (extract_query policy_code))
+
+let test_local_not_ground_because_of_ordering () =
+  let policy_code = {|
+package test
+
+q {
+  y := x
+  x := 2
+}
+  |} in
+  Alcotest.(check strings)
+    "expected identifier list"
+    ["x"]
+    (To_test.extract_locals (extract_query policy_code))
+
+
+let () =
+  let open Alcotest in
+  run "Compiler"
+    [ "local extraction", [
+        test_case "local extraction simple (1)"
+          `Quick test_local_extraction_simple_1;
+        test_case "local extraction simple (2)"
+          `Quick test_local_extraction_simple_2;
+        test_case "local extraction with local dependent on previously-defined local (1)"
+          `Quick test_local_dependent_on_previously_defined_local;
+        test_case "local extraction with local dependent on previously-defined local (2)"
+          `Quick test_local_not_ground_because_of_ordering ] ]