1 files changed, 26 insertions, 42 deletions

diff --git a/test/benchmarks/base/startup_test.go b/test/benchmarks/base/startup_test.go
index c36a544db..28731f97a 100644
--- a/test/benchmarks/base/startup_test.go
+++ b/test/benchmarks/base/startup_test.go
@@ -12,19 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package base
+package startup_test
 
 import (
 	"context"
 	"fmt"
-	"net"
+	"os"
 	"testing"
-	"time"
 
 	"gvisor.dev/gvisor/pkg/test/dockerutil"
+	"gvisor.dev/gvisor/test/benchmarks/base"
 	"gvisor.dev/gvisor/test/benchmarks/harness"
 )
 
+var testHarness harness.Harness
+
 // BenchmarkStartEmpty times startup time for an empty container.
 func BenchmarkStartupEmpty(b *testing.B) {
 	machine, err := testHarness.GetMachine()
@@ -60,11 +62,11 @@ func BenchmarkStartupNginx(b *testing.B) {
 		Image: "benchmarks/nginx",
 	}
 	runServerWorkload(ctx, b,
-		serverArgs{
-			machine: machine,
-			runOpts: runOpts,
-			port:    80,
-			cmd:     []string{"nginx", "-c", "/etc/nginx/nginx_gofer.conf"},
+		base.ServerArgs{
+			Machine: machine,
+			RunOpts: runOpts,
+			Port:    80,
+			Cmd:     []string{"nginx", "-c", "/etc/nginx/nginx_gofer.conf"},
 		})
 }
 
@@ -79,7 +81,7 @@ func BenchmarkStartupNode(b *testing.B) {
 	defer machine.CleanUp()
 
 	ctx := context.Background()
-	redis, redisIP := redisInstance(ctx, b, machine)
+	redis, redisIP := base.RedisInstance(ctx, b, machine)
 	defer redis.CleanUp(ctx)
 	runOpts := dockerutil.RunOpts{
 		Image:   "benchmarks/node",
@@ -89,52 +91,28 @@ func BenchmarkStartupNode(b *testing.B) {
 
 	cmd := []string{"node", "index.js", redisIP.String()}
 	runServerWorkload(ctx, b,
-		serverArgs{
-			machine: machine,
-			port:    8080,
-			runOpts: runOpts,
-			cmd:     cmd,
+		base.ServerArgs{
+			Machine: machine,
+			Port:    8080,
+			RunOpts: runOpts,
+			Cmd:     cmd,
 		})
 }
 
-// redisInstance returns a Redis container and its reachable IP.
-func redisInstance(ctx context.Context, b *testing.B, machine harness.Machine) (*dockerutil.Container, net.IP) {
-	b.Helper()
-	// Spawn a redis instance for the app to use.
-	redis := machine.GetNativeContainer(ctx, b)
-	if err := redis.Spawn(ctx, dockerutil.RunOpts{
-		Image: "benchmarks/redis",
-	}); err != nil {
-		redis.CleanUp(ctx)
-		b.Fatalf("failed to spwan redis instance: %v", err)
-	}
-
-	if out, err := redis.WaitForOutput(ctx, "Ready to accept connections", 3*time.Second); err != nil {
-		redis.CleanUp(ctx)
-		b.Fatalf("failed to start redis server: %v %s", err, out)
-	}
-	redisIP, err := redis.FindIP(ctx, false)
-	if err != nil {
-		redis.CleanUp(ctx)
-		b.Fatalf("failed to get IP from redis instance: %v", err)
-	}
-	return redis, redisIP
-}
-
 // runServerWorkload runs a server workload defined by 'runOpts' and 'cmd'.
 // 'clientMachine' is used to connect to the server on 'serverMachine'.
-func runServerWorkload(ctx context.Context, b *testing.B, args serverArgs) {
+func runServerWorkload(ctx context.Context, b *testing.B, args base.ServerArgs) {
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		if err := func() error {
-			server := args.machine.GetContainer(ctx, b)
+			server := args.Machine.GetContainer(ctx, b)
 			defer func() {
 				b.StopTimer()
 				// Cleanup servers as we run so that we can go indefinitely.
 				server.CleanUp(ctx)
 				b.StartTimer()
 			}()
-			if err := server.Spawn(ctx, args.runOpts, args.cmd...); err != nil {
+			if err := server.Spawn(ctx, args.RunOpts, args.Cmd...); err != nil {
 				return fmt.Errorf("failed to spawn node instance: %v", err)
 			}
 
@@ -144,7 +122,7 @@ func runServerWorkload(ctx context.Context, b *testing.B, args serverArgs) {
 			}
 
 			// Wait until the Client sees the server as up.
-			if err := harness.WaitUntilServing(ctx, args.machine, servingIP, args.port); err != nil {
+			if err := harness.WaitUntilServing(ctx, args.Machine, servingIP, args.Port); err != nil {
 				return fmt.Errorf("failed to wait for serving: %v", err)
 			}
 			return nil
@@ -153,3 +131,9 @@ func runServerWorkload(ctx context.Context, b *testing.B, args serverArgs) {
 		}
 	}
 }
+
+// TestMain is the main method for package network.
+func TestMain(m *testing.M) {
+	testHarness.Init()
+	os.Exit(m.Run())
+}