Move main methods for benchmark packages main package file.

PiperOrigin-RevId: 321875119

3 files changed, 16 insertions, 9 deletions

diff --git a/test/benchmarks/network/BUILD b/test/benchmarks/network/BUILD
index ea78416cf..16d267bc8 100644
--- a/test/benchmarks/network/BUILD
+++ b/test/benchmarks/network/BUILD
@@ -6,6 +6,7 @@ go_library(
     name = "network",
     testonly = 1,
     srcs = ["network.go"],
+    deps = ["//test/benchmarks/harness"],
 )
 
 go_test(
diff --git a/test/benchmarks/network/iperf_test.go b/test/benchmarks/network/iperf_test.go
index 48cc9dd8f..664e0797e 100644
--- a/test/benchmarks/network/iperf_test.go
+++ b/test/benchmarks/network/iperf_test.go
@@ -16,7 +16,6 @@ package network
 import (
 	"context"
 	"fmt"
-	"os"
 	"regexp"
 	"strconv"
 	"strings"
@@ -26,8 +25,6 @@ import (
 	"gvisor.dev/gvisor/test/benchmarks/harness"
 )
 
-var h harness.Harness
-
 func BenchmarkIperf(b *testing.B) {
 
 	// Get two machines
@@ -144,10 +141,4 @@ TCP window size: 45.0 KByte (default)
 	if err != nil || bandwidth != 45900 {
 		t.Fatalf("failed with: %v and %f", err, bandwidth)
 	}
-
-}
-
-func TestMain(m *testing.M) {
-	h.Init()
-	os.Exit(m.Run())
 }
diff --git a/test/benchmarks/network/network.go b/test/benchmarks/network/network.go
index f480b5bcd..ce17ddb94 100644
--- a/test/benchmarks/network/network.go
+++ b/test/benchmarks/network/network.go
@@ -14,3 +14,18 @@
 
 // Package network holds benchmarks around raw network performance.
 package network
+
+import (
+	"os"
+	"testing"
+
+	"gvisor.dev/gvisor/test/benchmarks/harness"
+)
+
+var h harness.Harness
+
+// TestMain is the main method for package network.
+func TestMain(m *testing.M) {
+	h.Init()
+	os.Exit(m.Run())
+}