wip: refactored the data model

Author: not-matthias

go-runner/src/builder/discovery.rs

@@ -1,7 +1,7 @@
 //! Finds all the benchmarks and packages in a given Go project.
 
 use std::{
-    collections::HashMap,
+    ops::Deref,
     path::{Path, PathBuf},
     process::Command,
 };
@@ -11,7 +11,7 @@ use crate::prelude::*;
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Module {
+pub struct GoModule {
     /// The module path (e.g., "local.dev/example-complex").
     #[serde(rename = "Path")]
     pub path: String,
@@ -58,7 +58,7 @@ pub struct GoPackage {
     pub compiled_go_files: Option<Vec<String>>,
 
     #[serde(rename = "Module")]
-    pub module: Module,
+    pub module: GoModule,
 }
 
 impl GoPackage {
@@ -70,15 +70,7 @@ impl GoPackage {
         serde_json::from_str(&format!("[{output}]")).context("Failed to parse Go list output")
     }
 
-    pub fn name(&self) -> &str {
-        &self.name
-    }
-
-    pub fn test_go_files(&self) -> Option<&Vec<String>> {
-        self.test_go_files.as_ref()
-    }
-
-    pub fn benchmarks(&self) -> anyhow::Result<Vec<BenchmarkData>> {
+    fn benchmarks(&self) -> anyhow::Result<Vec<GoBenchmark>> {
         let Some(test_go_files) = &self.test_go_files else {
             bail!("No test files found for package: {}", self.name);
         };
@@ -100,25 +92,25 @@ impl GoPackage {
             };
 
             // We don't support files with benchmarks that also contains github.com/frankban/quicktest
-            if file
-                .imports
-                .iter()
-                .any(|import| import.path.value.contains("github.com/frankban/quicktest"))
-            {
-                warn!("Skipping file with quicktest import: {file_path:?}");
-                continue;
-            }
-
-            // We don't support testify yet
-            // FIXME: We shouldn't skip packages where benchmarks and tests are separated -> opentelemetry-go
-            if file
-                .imports
-                .iter()
-                .any(|import| import.path.value.contains("github.com/stretchr/testify"))
-            {
-                warn!("Skipping file with testify import: {file_path:?}");
-                continue;
-            }
+            // if file
+            //     .imports
+            //     .iter()
+            //     .any(|import| import.path.value.contains("github.com/frankban/quicktest"))
+            // {
+            //     warn!("Skipping file with quicktest import: {file_path:?}");
+            //     continue;
+            // }
+
+            // // We don't support testify yet
+            // // FIXME: We shouldn't skip packages where benchmarks and tests are separated -> opentelemetry-go
+            // if file
+            //     .imports
+            //     .iter()
+            //     .any(|import| import.path.value.contains("github.com/stretchr/testify"))
+            // {
+            //     warn!("Skipping file with testify import: {file_path:?}");
+            //     continue;
+            // }
 
             // We can't import packages that are declared as `main`
             if file.pkg_name.name == "main" {
@@ -159,8 +151,26 @@ impl GoPackage {
 
             let valid_benchmarks =
                 verifier::FuncVisitor::verify_source_code(&content, &found_benchmarks)?;
+            if valid_benchmarks.len() != found_benchmarks.len() {
+                warn!(
+                    "Only {} out of {} are valid, skipping file",
+                    valid_benchmarks.len(),
+                    found_benchmarks.len()
+                );
+                warn!("Valid benchmarks: {valid_benchmarks:?}");
+                warn!(
+                    "Invalid benchmarks: {:?}",
+                    found_benchmarks
+                        .iter()
+                        .filter(|f| !valid_benchmarks.contains(f))
+                        .collect::<Vec<_>>()
+                );
+
+                continue;
+            }
+
             for func in valid_benchmarks {
-                benchmarks.push(BenchmarkData::new(
+                benchmarks.push(GoBenchmark::new(
                     package_import_path.clone(),
                     func,
                     root_relative_file_path.to_path_buf(),
@@ -173,7 +183,7 @@ impl GoPackage {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct BenchmarkData {
+pub struct GoBenchmark {
     /// The name of the benchmark (e.g. `BenchmarkFoo`).
     pub name: String,
 
@@ -190,7 +200,7 @@ pub struct BenchmarkData {
     pub file_path: PathBuf,
 }
 
-impl BenchmarkData {
+impl GoBenchmark {
     pub fn new(package_import_path: String, name: String, file_path: PathBuf) -> Self {
         #[cfg(test)]
         let import_alias = {
@@ -213,143 +223,129 @@ impl BenchmarkData {
     }
 }
 
-fn run_go_list(go_project_path: &Path) -> anyhow::Result<Vec<GoPackage>> {
-    // Execute 'go list -test -compiled -json ./...' to get package information
-    let output = Command::new("go")
-        .args(["list", "-test", "-compiled", "-json", "./..."])
-        .current_dir(go_project_path)
-        .output()?;
-
-    if !output.status.success() {
-        bail!(
-            "Failed to execute 'go list': {}",
-            String::from_utf8_lossy(&output.stderr)
-        );
+/// Represents a package with its benchmarks.
+#[derive(Debug, Clone, Serialize)]
+pub struct BenchmarkPackage {
+    raw_package: GoPackage,
+    pub benchmarks: Vec<GoBenchmark>,
+}
+
+impl BenchmarkPackage {
+    fn new(package: GoPackage, benchmarks: Vec<GoBenchmark>) -> Self {
+        Self {
+            raw_package: package,
+            benchmarks,
+        }
     }
 
-    // Wrap it in '[{output}]' and parse it with serde_json
-    let output_str = String::from_utf8(output.stdout)?;
-    trace!("Go list output: {output_str}");
+    pub fn from_project(go_project_path: &Path) -> anyhow::Result<Vec<BenchmarkPackage>> {
+        let raw_packages = Self::run_go_list(go_project_path)?;
+        let has_test_files =
+            |files: &Vec<String>| files.iter().any(|name| name.ends_with("_test.go"));
+        let has_test_imports = |imports: &Vec<String>| {
+            imports.iter().any(|import| {
+                // import "testing"
+                import.contains("testing")
+            })
+        };
 
-    GoPackage::from_go_list_output(&output_str)
-}
+        let mut packages = Vec::new();
+        for package in raw_packages {
+            // Skip packages without test files
+            let has_tests = package
+                .test_go_files
+                .as_ref()
+                .map(has_test_files)
+                .unwrap_or_default();
+            if !has_tests {
+                debug!("Skipping package without test files: {}", package.name);
+                continue;
+            }
 
-pub fn run(go_project_path: &Path) -> anyhow::Result<HashMap<String, GoPackage>> {
-    let packages = run_go_list(go_project_path)?;
+            // Skip packages without test imports
+            let has_test_imports = package
+                .test_imports
+                .as_ref()
+                .map(has_test_imports)
+                .unwrap_or_default();
+            if !has_test_imports {
+                debug!("Skipping package without test imports: {}", package.name);
+                continue;
+            }
 
-    discover_benchmarks(packages)
-}
+            // Only include test executables, since we want to generate them manually.
+            // Example format: `local.dev/example-complex [local.dev/example-complex.test]`
+            if !package.import_path.ends_with(".test]") {
+                debug!("Skipping package without test executable: {}", package.name);
+                continue;
+            }
 
-fn discover_benchmarks(packages: Vec<GoPackage>) -> anyhow::Result<HashMap<String, GoPackage>> {
-    let has_test_files = |files: &Vec<String>| files.iter().any(|name| name.ends_with("_test.go"));
-    let has_test_imports = |imports: &Vec<String>| {
-        imports.iter().any(|import| {
-            // import "testing"
-            import.contains("testing")
-        })
-    };
-
-    let mut benchmark_packages: HashMap<String, GoPackage> = HashMap::new();
-    for package in packages {
-        // Skip packages without test files
-        let has_tests = package
-            .test_go_files
-            .as_ref()
-            .map(has_test_files)
-            .unwrap_or_default();
-        if !has_tests {
-            debug!("Skipping package without test files: {}", package.name);
-            continue;
-        }
+            // Skip packages that don't have benchmarks
+            let benchmarks = match package.benchmarks() {
+                Ok(benchmarks) => benchmarks,
+                Err(e) => {
+                    warn!(
+                        "Failed to get benchmarks for package {}: {}",
+                        package.name, e
+                    );
+                    continue;
+                }
+            };
+            if benchmarks.is_empty() {
+                debug!("Skipping package without benchmarks: {}", package.name);
+                continue;
+            }
 
-        // Skip packages without test imports
-        let has_test_imports = package
-            .test_imports
-            .as_ref()
-            .map(has_test_imports)
-            .unwrap_or_default();
-        if !has_test_imports {
-            debug!("Skipping package without test imports: {}", package.name);
-            continue;
+            packages.push(BenchmarkPackage::new(package, benchmarks));
         }
 
-        // Only include test executables, since we want to generate them manually.
-        // Example format: `local.dev/example-complex [local.dev/example-complex.test]`
-        if !package.import_path.ends_with(".test]") {
-            debug!("Skipping package without test executable: {}", package.name);
-            continue;
-        }
+        Ok(packages)
+    }
 
-        // Skip packages that don't have benchmarks
-        let benchmarks = match package.benchmarks() {
-            Ok(benchmarks) => benchmarks,
-            Err(e) => {
-                warn!(
-                    "Failed to get benchmarks for package {}: {}",
-                    package.name, e
-                );
-                continue;
-            }
-        };
-        if benchmarks.is_empty() {
-            debug!("Skipping package without benchmarks: {}", package.name);
-            continue;
+    fn run_go_list(go_project_path: &Path) -> anyhow::Result<Vec<GoPackage>> {
+        // Execute 'go list -test -compiled -json ./...' to get package information
+        let output = Command::new("go")
+            .args(["list", "-test", "-compiled", "-json", "./..."])
+            .current_dir(go_project_path)
+            .output()?;
+
+        if !output.status.success() {
+            bail!(
+                "Failed to execute 'go list': {}",
+                String::from_utf8_lossy(&output.stderr)
+            );
         }
 
-        assert!(!benchmark_packages.contains_key(&package.import_path));
-        benchmark_packages.insert(package.import_path.clone(), package);
-        // // FIXME: Assert that we didn't already add this package
-        // benchmark_packages
-        //     .entry(package.import_path.clone())
-        //     .insert_entry(package);
+        // Wrap it in '[{output}]' and parse it with serde_json
+        let output_str = String::from_utf8(output.stdout)?;
+        trace!("Go list output: {output_str}");
+
+        GoPackage::from_go_list_output(&output_str)
     }
+}
 
-    Ok(benchmark_packages)
+impl Deref for BenchmarkPackage {
+    type Target = GoPackage;
+
+    fn deref(&self) -> &Self::Target {
+        &self.raw_package
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    #[test]
-    fn test_parse_go_list_output() {
-        let packages = run_go_list(Path::new("testdata/projects/golang-benchmarks")).unwrap();
-
-        insta::with_settings!({sort_maps => true}, {
-            insta::assert_json_snapshot!(packages, {
-                ".**[\"Dir\"]" => "[package_dir]",
-                ".**[\"Module\"][\"Dir\"]" => "[module_dir]",
-                ".**[\"Module\"][\"GoMod\"]" => "[go_mod_path]"
-            });
-        });
-    }
-
     #[test]
     fn test_discover_benchmarks() {
-        let packages = run_go_list(Path::new("testdata/projects/golang-benchmarks")).unwrap();
-        let packages = discover_benchmarks(packages).unwrap();
-
-        insta::with_settings!({sort_maps => true}, {
-            insta::assert_json_snapshot!(packages, {
-                ".**[\"Dir\"]" => "[package_dir]",
-                ".**[\"Module\"][\"Dir\"]" => "[module_dir]",
-                ".**[\"Module\"][\"GoMod\"]" => "[go_mod_path]"
-            });
+        let packages =
+            BenchmarkPackage::from_project(Path::new("testdata/projects/golang-benchmarks"))
+                .unwrap();
+
+        insta::assert_json_snapshot!(packages, {
+            ".**[\"Dir\"]" => "[package_dir]",
+            ".**[\"Module\"][\"Dir\"]" => "[module_dir]",
+            ".**[\"Module\"][\"GoMod\"]" => "[go_mod_path]"
         });
     }
-
-    #[test]
-    fn test_benchmarks_for_package() {
-        let packages = run_go_list(Path::new("testdata/projects/golang-benchmarks")).unwrap();
-        let benchmark_packages = discover_benchmarks(packages).unwrap();
-
-        let package = benchmark_packages
-            .get("github.com/SimonWaldherr/golang-benchmarks/base64 [github.com/SimonWaldherr/golang-benchmarks/base64.test]")
-            .unwrap();
-
-        let mut benches = package.benchmarks().unwrap();
-        benches.sort_by_cached_key(|b| b.name.clone());
-
-        insta::assert_json_snapshot!(benches);
-    }
 }