From 043ce9c5d2b0ee7d5d476c5c8475cd44b724027a Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Fri, 16 Aug 2024 20:51:51 -0700
Subject: [PATCH] Fix `sniffer_test` by embedding the `run_sniffer` binary in
 it.

This works around file path resolution issues.

PiperOrigin-RevId: 663985469
---
 Makefile                 |  4 ++--
 test/gpu/BUILD           | 16 +++++++++++++---
 test/gpu/sniffer_test.go | 33 ++++++++++++++++++++++++++++-----
 3 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/Makefile b/Makefile
index 1f35f73171..bb22627816 100644
--- a/Makefile
+++ b/Makefile
@@ -308,7 +308,7 @@ gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:nccl_test,--runtime=$(RUNTIME) -test.v $(ARGS))
-	@$(call test,--test_env=RUNTIME=$(RUNTIME) test/gpu:sniffer_test)
+	@$(call sudo,test/gpu:sniffer_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-all-tests
 
 cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
@@ -318,7 +318,7 @@ cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:nccl_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
-	@$(call test,--test_env=RUNTIME=$(RUNTIME) --test_arg="--cos-gpu" test/gpu:sniffer_test)
+	@$(call sudo,test/gpu:sniffer_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-all-tests
 
 portforward-tests: load-basic_redis load-basic_nginx $(RUNTIME_BIN)
diff --git a/test/gpu/BUILD b/test/gpu/BUILD
index 89d1875dda..0fd44e2cc7 100644
--- a/test/gpu/BUILD
+++ b/test/gpu/BUILD
@@ -103,11 +103,22 @@ go_test(
     deps = ["//test/gpu/stablediffusion"],
 )
 
+# We copy the `run_sniffer` binary here because `go:embed` can only embed
+# from the current directory or subdirectories, not parents of it.
+genrule(
+    name = "run_sniffer_copy",
+    srcs = [
+        "//tools/ioctl_sniffer:run_sniffer",
+    ],
+    outs = ["run_sniffer_copy"],
+    cmd = "cat < $(SRCS) > $@",
+)
+
 go_test(
     name = "sniffer_test",
     srcs = ["sniffer_test.go"],
-    data = [
-        "//tools/ioctl_sniffer:run_sniffer",
+    embedsrcs = [
+        ":run_sniffer_copy",  # keep
     ],
     tags = [
         "manual",
@@ -117,7 +128,6 @@ go_test(
     visibility = ["//:sandbox"],
     deps = [
         "//pkg/test/dockerutil",
-        "//pkg/test/testutil",
         "@com_github_docker_docker//api/types/mount:go_default_library",
     ],
 )
diff --git a/test/gpu/sniffer_test.go b/test/gpu/sniffer_test.go
index f10debdbca..6d7e088af2 100644
--- a/test/gpu/sniffer_test.go
+++ b/test/gpu/sniffer_test.go
@@ -18,25 +18,48 @@ package sniffer_test
 import (
 	"context"
 	"errors"
+	"os"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/docker/docker/api/types/mount"
 	"gvisor.dev/gvisor/pkg/test/dockerutil"
-	"gvisor.dev/gvisor/pkg/test/testutil"
+
+	// Needed for go:embed
+	_ "embed"
 )
 
 const maxDuration = 1 * time.Minute
 
+//go:embed run_sniffer_copy
+var runSnifferBinary []byte
+
 // RunCommand runs the given command via the sniffer, with the -enforce_compatibility flag.
 //
 //	It's run in a docker container, with the cuda-tests image.
 func runCUDATestsCommand(t *testing.T, cmd ...string) (string, error) {
-	// Find the sniffer binary
-	cliPath, err := testutil.FindFile("tools/ioctl_sniffer/run_sniffer")
+	// Extract the sniffer binary to a temporary location.
+	runSniffer, err := os.CreateTemp("/tmp", "run_sniffer.*")
 	if err != nil {
-		t.Fatalf("Failed to find run_sniffer: %v", err)
+		t.Fatalf("Failed to create temporary file: %v", err)
+	}
+	defer func() {
+		if err := runSniffer.Close(); err != nil {
+			t.Fatalf("Failed to close temporary file: %v", err)
+		}
+		if err := os.Remove(runSniffer.Name()); err != nil {
+			t.Fatalf("Failed to unlink temporary file: %v", err)
+		}
+	}()
+	if _, err := runSniffer.Write(runSnifferBinary); err != nil {
+		t.Fatalf("Failed to write to temporary file: %v", err)
+	}
+	if err := runSniffer.Sync(); err != nil {
+		t.Fatalf("Failed to sync temporary file: %v", err)
+	}
+	if err := runSniffer.Chmod(0o555); err != nil {
+		t.Fatalf("Failed to chmod temporary file: %v", err)
 	}
 
 	// Set up our docker container
@@ -51,7 +74,7 @@ func runCUDATestsCommand(t *testing.T, cmd ...string) (string, error) {
 	opts.Image = "gpu/cuda-tests"
 	opts.Mounts = append(opts.Mounts, mount.Mount{
 		Type:     mount.TypeBind,
-		Source:   cliPath,
+		Source:   runSniffer.Name(),
 		Target:   "/run_sniffer",
 		ReadOnly: false,
 	})