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))
+    }
+    }
 }