feat: suport tests in main package

go-runner/Cargo.toml

@@ -24,3 +24,4 @@ tempfile = "3.14"
 [dev-dependencies]
 insta = { version = "1.43", features = ["json", "redactions"] }
 rstest = "0.26"
+test-log = "0.2.18"

go-runner/src/builder/discovery.rs

@@ -93,12 +93,6 @@ impl GoPackage {
                 }
             };
 
-            // We can't import packages that are declared as `main`
-            if file.pkg_name.name == "main" {
-                warn!("Skipping file with main package: {file_path:?}");
-                continue;
-            }
-
             // First, collect all benchmark function names from this file
             let mut found_benchmarks = Vec::new();
             for decl in &file.decl {

go-runner/src/builder/patcher.rs

@@ -114,8 +114,9 @@ pub fn patch_go_source(source: &str) -> anyhow::Result<String> {
             Ok(source)
         };
 
+    let source = replace_package_main(source.into())?;
     let source = replace_import(
-        source.to_string(),
+        source,
         "testing",
         "testing \"github.com/CodSpeedHQ/codspeed-go/compat/testing\"",
     )?;
@@ -133,6 +134,24 @@ pub fn patch_go_source(source: &str) -> anyhow::Result<String> {
     Ok(source)
 }
 
+/// Replace `package main` with `package main_compat` to allow importing it from other packages.
+fn replace_package_main(source: String) -> anyhow::Result<String> {
+    let parsed = gosyn::parse_source(&source)?;
+
+    // Only replace if package name is "main"
+    if parsed.pkg_name.name != "main" {
+        return Ok(source);
+    }
+
+    // pkg_name.pos is the position of the identifier "main" in the source
+    let name_start = parsed.pkg_name.pos;
+    let name_end = name_start + parsed.pkg_name.name.len();
+
+    let mut result = source;
+    result.replace_range(name_start..name_end, "main_compat");
+    Ok(result)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -244,6 +263,19 @@ import (
     "fmt"
     "github.com/thejerf/slogassert"
 )
+"#;
+
+    const PACKAGE_MAIN: &str = r#"package main
+
+import "testing"
+
+func BenchmarkExample(b *testing.B) {
+    // benchmark code
+}
+
+func TestExample(t *testing.T) {
+    s := "package main"
+}
 "#;
 
     #[rstest]
@@ -261,6 +293,7 @@ import (
         "multiline_import_with_testing_string",
         MULTILINE_IMPORT_WITH_TESTING_STRING
     )]
+    #[case("package_main", PACKAGE_MAIN)]
     fn test_patch_go_source(#[case] test_name: &str, #[case] source: &str) {
         let result = patch_go_source(source).unwrap();
         assert_snapshot!(test_name, result);

go-runner/src/builder/snapshots/codspeed_go_runner__builder__discovery__tests__discover_benchmarks.snap

@@ -342,6 +342,59 @@ expression: packages
       }
     ]
   },
+  {
+    "raw_package": {
+      "Dir": "[package_dir]",
+      "Name": "main",
+      "ImportPath": "github.com/SimonWaldherr/golang-benchmarks/floodfill [github.com/SimonWaldherr/golang-benchmarks/floodfill.test]",
+      "TestGoFiles": [
+        "floodfill_test.go"
+      ],
+      "TestImports": [
+        "testing"
+      ],
+      "CompiledGoFiles": [
+        "floodfill_test.go"
+      ],
+      "Module": {
+        "Path": "github.com/SimonWaldherr/golang-benchmarks",
+        "Dir": "[module_dir]",
+        "GoMod": "[go_mod_path]",
+        "GoVersion": "1.18",
+        "Main": true
+      }
+    },
+    "benchmarks": [
+      {
+        "name": "BenchmarkFloodFillBFS",
+        "module_path": "github.com/SimonWaldherr/golang-benchmarks/floodfill",
+        "import_alias": "benchmarkfloodfillbfs_18158933623445949257",
+        "qualified_name": "benchmarkfloodfillbfs_18158933623445949257.BenchmarkFloodFillBFS",
+        "file_path": "floodfill/floodfill_test.go"
+      },
+      {
+        "name": "BenchmarkFloodFillDFS",
+        "module_path": "github.com/SimonWaldherr/golang-benchmarks/floodfill",
+        "import_alias": "benchmarkfloodfilldfs_18158933623445949257",
+        "qualified_name": "benchmarkfloodfilldfs_18158933623445949257.BenchmarkFloodFillDFS",
+        "file_path": "floodfill/floodfill_test.go"
+      },
+      {
+        "name": "BenchmarkFloodFillRecursive",
+        "module_path": "github.com/SimonWaldherr/golang-benchmarks/floodfill",
+        "import_alias": "benchmarkfloodfillrecursive_18158933623445949257",
+        "qualified_name": "benchmarkfloodfillrecursive_18158933623445949257.BenchmarkFloodFillRecursive",
+        "file_path": "floodfill/floodfill_test.go"
+      },
+      {
+        "name": "BenchmarkFloodFillStack4Way",
+        "module_path": "github.com/SimonWaldherr/golang-benchmarks/floodfill",
+        "import_alias": "benchmarkfloodfillstack4way_18158933623445949257",
+        "qualified_name": "benchmarkfloodfillstack4way_18158933623445949257.BenchmarkFloodFillStack4Way",
+        "file_path": "floodfill/floodfill_test.go"
+      }
+    ]
+  },
   {
     "raw_package": {
       "Dir": "[package_dir]",

go-runner/src/builder/snapshots/codspeed_go_runner__builder__patcher__tests__already_patched_import.snap

@@ -1,8 +1,8 @@
 ---
-source: src/builder/patcher.rs
+source: go-runner/src/builder/patcher.rs
 expression: result
 ---
-package main
+package main_compat
 
 import testing "github.com/CodSpeedHQ/codspeed-go/compat/testing"
 

go-runner/src/builder/snapshots/codspeed_go_runner__builder__patcher__tests__import_at_end_of_block.snap

@@ -1,8 +1,8 @@
 ---
-source: src/builder/patcher.rs
+source: go-runner/src/builder/patcher.rs
 expression: result
 ---
-package main
+package main_compat
 
 import (
     "fmt"

go-runner/src/builder/snapshots/codspeed_go_runner__builder__patcher__tests__import_testing_and_slogassert.snap

@@ -2,7 +2,7 @@
 source: go-runner/src/builder/patcher.rs
 expression: result
 ---
-package main
+package main_compat
 import (
     testing "github.com/CodSpeedHQ/codspeed-go/compat/testing"
     "fmt"

go-runner/src/builder/snapshots/codspeed_go_runner__builder__patcher__tests__import_with_comments.snap

@@ -1,8 +1,8 @@
 ---
-source: src/builder/patcher.rs
+source: go-runner/src/builder/patcher.rs
 expression: result
 ---
-package main
+package main_compat
 
 import (
     "fmt"

go-runner/src/builder/snapshots/codspeed_go_runner__builder__patcher__tests__import_with_extra_whitespace.snap

@@ -1,8 +1,8 @@
 ---
-source: src/builder/patcher.rs
+source: go-runner/src/builder/patcher.rs
 expression: result
 ---
-package main
+package main_compat
 
 import (
     "fmt"