3 files changed, 18 insertions, 25 deletions
diff --git a/test/benchmarks/media/BUILD b/test/benchmarks/media/BUILD
index bb242d385..380783f0b 100644
--- a/test/benchmarks/media/BUILD
+++ b/test/benchmarks/media/BUILD
@@ -1,4 +1,5 @@
-load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools:defs.bzl", "go_library")
+load("//test/benchmarks:defs.bzl", "benchmark_test")
 
 package(licenses = ["notice"])
 
@@ -6,12 +7,11 @@ go_library(
     name = "media",
     testonly = 1,
     srcs = ["media.go"],
-    deps = ["//test/benchmarks/harness"],
 )
 
-go_test(
-    name = "media_test",
-    size = "large",
+benchmark_test(
+    name = "ffmpeg_test",
+    size = "enormous",
     srcs = ["ffmpeg_test.go"],
     library = ":media",
     visibility = ["//:sandbox"],
diff --git a/test/benchmarks/media/ffmpeg_test.go b/test/benchmarks/media/ffmpeg_test.go
index 7822dfad7..1b99a319a 100644
--- a/test/benchmarks/media/ffmpeg_test.go
+++ b/test/benchmarks/media/ffmpeg_test.go
@@ -15,6 +15,7 @@ package media
 
 import (
 	"context"
+	"os"
 	"strings"
 	"testing"
 
@@ -25,29 +26,36 @@ import (
 // BenchmarkFfmpeg runs ffmpeg in a container and records runtime.
 // BenchmarkFfmpeg should run as root to drop caches.
 func BenchmarkFfmpeg(b *testing.B) {
-	machine, err := h.GetMachine()
+	machine, err := harness.GetMachine()
 	if err != nil {
 		b.Fatalf("failed to get machine: %v", err)
 	}
 	defer machine.CleanUp()
 
 	ctx := context.Background()
-	container := machine.GetContainer(ctx, b)
-	defer container.CleanUp(ctx)
 	cmd := strings.Split("ffmpeg -i video.mp4 -c:v libx264 -preset veryslow output.mp4", " ")
 
 	b.ResetTimer()
+	b.StopTimer()
+
 	for i := 0; i < b.N; i++ {
-		b.StopTimer()
+		container := machine.GetContainer(ctx, b)
+		defer container.CleanUp(ctx)
 		if err := harness.DropCaches(machine); err != nil {
 			b.Skipf("failed to drop caches: %v. You probably need root.", err)
 		}
-		b.StartTimer()
 
+		b.StartTimer()
 		if _, err := container.Run(ctx, dockerutil.RunOpts{
 			Image: "benchmarks/ffmpeg",
 		}, cmd...); err != nil {
 			b.Fatalf("failed to run container: %v", err)
 		}
+		b.StopTimer()
 	}
 }
+
+func TestMain(m *testing.M) {
+	harness.Init()
+	os.Exit(m.Run())
+}
diff --git a/test/benchmarks/media/media.go b/test/benchmarks/media/media.go
index c7b35b758..ed7b24651 100644
--- a/test/benchmarks/media/media.go
+++ b/test/benchmarks/media/media.go
@@ -14,18 +14,3 @@
 
 // Package media holds benchmarks around media processing applications.
 package media
-
-import (
-	"os"
-	"testing"
-
-	"gvisor.dev/gvisor/test/benchmarks/harness"
-)
-
-var h harness.Harness
-
-// TestMain is the main method for package media.
-func TestMain(m *testing.M) {
-	h.Init()
-	os.Exit(m.Run())
-}