Initial implementation of policy coverage (#146)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

tests/aci/main.rs

@@ -108,6 +108,70 @@ fn run_aci_tests(dir: &Path) -> Result<()> {
     Ok(())
 }
 
+#[cfg(feature = "coverage")]
+fn run_aci_tests_coverage(dir: &Path) -> Result<()> {
+    let mut engine = Engine::new();
+
+    let mut added = std::collections::BTreeSet::new();
+
+    for entry in WalkDir::new(dir)
+        .sort_by_file_name()
+        .into_iter()
+        .filter_map(|e| e.ok())
+    {
+        let path = entry.path();
+        if !path.to_string_lossy().ends_with(".yaml") {
+            continue;
+        }
+
+        let yaml = std::fs::read(&path)?;
+        let yaml = String::from_utf8_lossy(&yaml);
+        let test: YamlTest = serde_yaml::from_str(&yaml)?;
+
+        for case in &test.cases {
+            for (idx, rego) in case.modules.iter().enumerate() {
+                if rego.ends_with(".rego") {
+                    let path = dir.join(rego);
+                    let path = path.to_str().expect("not a valid path");
+                    let path = path.to_string();
+                    if !added.contains(&path) {
+                        engine.add_policy_from_file(path.to_string())?;
+                        added.insert(path);
+                    }
+                } else {
+                    engine.add_policy(format!("rego{idx}.rego"), rego.clone())?;
+                }
+            }
+
+            engine.clear_data();
+            engine.add_data(case.data.clone())?;
+            engine.set_input(case.input.clone());
+            let _query_results = engine.eval_query(case.query.clone(), true)?;
+        }
+    }
+
+    println!("\n\nCOVERAGE REPORT");
+    // Fetch coverage report.
+    let report = engine.get_coverage_report()?;
+    for file in report.files.into_iter() {
+        if file.uncovered.is_empty() {
+            println!("{} has full coverage", file.path);
+            continue;
+        }
+
+        println!("{}:", file.path);
+        for (line, code) in file.code.split('\n').enumerate() {
+            if file.uncovered.contains(&(line as u32 + 1)) {
+                println!("\x1b[31m {line:4}  {code}\x1b[0m");
+            } else {
+                println!(" {line:4}  {code}");
+            }
+        }
+    }
+
+    Ok(())
+}
+
 #[derive(clap::Parser)]
 #[command(author, version, about, long_about = None)]
 struct Cli {
@@ -119,5 +183,12 @@ struct Cli {
 
 fn main() -> Result<()> {
     let cli = Cli::parse();
-    run_aci_tests(&Path::new(&cli.test_dir))
+
+    cfg_if::cfg_if! {
+    if #[cfg(feature = "coverage")] {
+        run_aci_tests_coverage(&Path::new(&cli.test_dir))
+    } else {
+        run_aci_tests(&Path::new(&cli.test_dir))
+    }
+    }
 }

tests/coverage/mod.rs

@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use std::collections::BTreeSet;
+
+use regorus::*;
+
+use anyhow::Result;
+use test_generator::test_resources;
+
+#[derive(serde::Deserialize)]
+struct TestCase {
+    data: Option<Value>,
+    input: Option<Value>,
+    modules: Vec<String>,
+    note: String,
+    query: String,
+    uncovered: Vec<BTreeSet<u32>>,
+    skip: Option<bool>,
+}
+
+#[derive(serde::Deserialize)]
+struct YamlTest {
+    cases: Vec<TestCase>,
+}
+
+fn yaml_test_impl(file: &str) -> Result<()> {
+    let yaml_str = std::fs::read_to_string(file)?;
+    let test: YamlTest = serde_yaml::from_str(&yaml_str)?;
+
+    println!("running {file}");
+
+    for case in test.cases.into_iter() {
+        print!("case {} ", case.note);
+        if case.skip == Some(true) {
+            println!("skipped");
+            continue;
+        }
+
+        let mut engine = Engine::new();
+
+        for (idx, rego) in case.modules.iter().enumerate() {
+            engine.add_policy(format!("rego_{idx}"), rego.clone())?;
+        }
+
+        if let Some(data) = case.data {
+            engine.add_data(data)?;
+        }
+
+        if let Some(input) = case.input {
+            engine.set_input(input);
+        }
+
+        let _ = engine.eval_query(case.query.clone(), false)?;
+        let report = engine.get_coverage_report()?;
+
+        for (idx, uncovered) in case.uncovered.into_iter().enumerate() {
+            assert_eq!(uncovered, report.files[idx].uncovered);
+        }
+
+        println!("passed");
+    }
+
+    Ok(())
+}
+
+fn yaml_test(file: &str) -> Result<()> {
+    match yaml_test_impl(file) {
+        Ok(_) => Ok(()),
+        Err(e) => {
+            // If Err is returned, it doesn't always get printed by cargo test.
+            // Therefore, panic with the error.
+            panic!("{}", e);
+        }
+    }
+}
+
+#[test_resources("tests/coverage/*.yaml")]
+fn run(path: &str) {
+    yaml_test(path).unwrap()
+}

tests/coverage/tests.yaml

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: basic
+    modules:
+      - |
+        package test
+
+        x = 1
+        
+        y = k {
+          input.x == 5
+          k = input.k
+        }
+    query: data.test
+    uncovered: [
+      [ 5, 7 ]
+    ]
+        

tests/mod.rs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+#[cfg(feature = "coverage")]
+mod coverage;
+
 mod engine;
 mod lexer;
 mod parser;