diff --git a/java/droidstubs_test.go b/java/droidstubs_test.go
index f86e1ac..7b89dde 100644
--- a/java/droidstubs_test.go
+++ b/java/droidstubs_test.go
@@ -84,7 +84,7 @@
 	for _, c := range testcases {
 		m := ctx.ModuleForTests(c.moduleName, "android_common")
 		manifest := m.Output("metalava.sbox.textproto")
-		sboxProto := android.RuleBuilderSboxProtoForTests(t, manifest)
+		sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx, manifest)
 		cmdline := String(sboxProto.Commands[0].Command)
 		android.AssertStringContainsEquals(t, "api-versions generation flag", cmdline, "--generate-api-levels", c.generate_xml)
 		if c.expectedJarFilename != "" {
@@ -133,7 +133,7 @@
 
 	m := ctx.ModuleForTests("foo-stubs", "android_common")
 	manifest := m.Output("metalava.sbox.textproto")
-	cmd := String(android.RuleBuilderSboxProtoForTests(t, manifest).Commands[0].Command)
+	cmd := String(android.RuleBuilderSboxProtoForTests(t, ctx, manifest).Commands[0].Command)
 	r := regexp.MustCompile(`--android-jar-pattern [^ ]+/android.jar`)
 	return r.FindAllString(cmd, -1)
 }
@@ -212,7 +212,7 @@
 		t.Errorf("Expected inputs %q, got %q", w, g)
 	}
 
-	manifest := android.RuleBuilderSboxProtoForTests(t, m.Output("metalava.sbox.textproto"))
+	manifest := android.RuleBuilderSboxProtoForTests(t, ctx, m.Output("metalava.sbox.textproto"))
 	if g, w := manifest.Commands[0].GetCommand(), "reference __SBOX_SANDBOX_DIR__/out/.intermediates/foo/gen/foo.txt"; !strings.Contains(g, w) {
 		t.Errorf("Expected command to contain %q, got %q", w, g)
 	}
@@ -302,7 +302,7 @@
 		})
 	m := ctx.ModuleForTests("baz-stubs", "android_common")
 	manifest := m.Output("metalava.sbox.textproto")
-	cmdline := String(android.RuleBuilderSboxProtoForTests(t, manifest).Commands[0].Command)
+	cmdline := String(android.RuleBuilderSboxProtoForTests(t, ctx, manifest).Commands[0].Command)
 	android.AssertStringDoesContain(t, "sdk-extensions-root present", cmdline, "--sdk-extensions-root sdk/extensions")
 	android.AssertStringDoesContain(t, "sdk-extensions-info present", cmdline, "--sdk-extensions-info sdk/extensions/info.txt")
 }
@@ -434,6 +434,6 @@
 
 	m := result.ModuleForTests("foo", "android_common")
 	manifest := m.Output("metalava.sbox.textproto")
-	cmdline := String(android.RuleBuilderSboxProtoForTests(t, manifest).Commands[0].Command)
+	cmdline := String(android.RuleBuilderSboxProtoForTests(t, result.TestContext, manifest).Commands[0].Command)
 	android.AssertStringDoesContain(t, "flagged api hide command not included", cmdline, "--hide-annotation android.annotation.FlaggedApi")
 }
diff --git a/java/java_test.go b/java/java_test.go
index 81119a7..0c750eb 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -1921,7 +1921,7 @@
 	for _, c := range testcases {
 		m := ctx.ModuleForTests(c.moduleName, "android_common")
 		manifest := m.Output("metalava.sbox.textproto")
-		sboxProto := android.RuleBuilderSboxProtoForTests(t, manifest)
+		sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx.TestContext, manifest)
 		manifestCommand := sboxProto.Commands[0].GetCommand()
 		sourceFilesFlag := "--source-files " + strings.Join(c.sourceTextFileDirs, " ")
 		android.AssertStringDoesContain(t, "source text files not present", manifestCommand, sourceFilesFlag)
@@ -2026,7 +2026,7 @@
 	for _, c := range testcases {
 		m := ctx.ModuleForTests(c.moduleName, "android_common")
 		manifest := m.Output("metalava.sbox.textproto")
-		sboxProto := android.RuleBuilderSboxProtoForTests(t, manifest)
+		sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx.TestContext, manifest)
 		manifestCommand := sboxProto.Commands[0].GetCommand()
 		sourceFilesFlag := "--source-files " + strings.Join(c.sourceTextFileDirs, " ")
 		android.AssertStringDoesContain(t, "source text files not present", manifestCommand, sourceFilesFlag)
@@ -2316,7 +2316,7 @@
 
 	m := ctx.ModuleForTests("bar1", "android_common")
 	manifest := m.Output("metalava.sbox.textproto")
-	sboxProto := android.RuleBuilderSboxProtoForTests(t, manifest)
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx.TestContext, manifest)
 	manifestCommand := sboxProto.Commands[0].GetCommand()
 	android.AssertStringDoesContain(t, "Command expected to contain full_api_surface_stub output jar", manifestCommand, "lib1.jar")
 }
@@ -2480,7 +2480,7 @@
 	`)
 	m := ctx.ModuleForTests("foo", "android_common")
 	manifest := m.Output("metalava.sbox.textproto")
-	sboxProto := android.RuleBuilderSboxProtoForTests(t, manifest)
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx.TestContext, manifest)
 	manifestCommand := sboxProto.Commands[0].GetCommand()
 	sourceFilesFlag := "--source-files current.txt"
 	android.AssertStringDoesContain(t, "source text files not present", manifestCommand, sourceFilesFlag)
@@ -2501,7 +2501,7 @@
 	`)
 	m := ctx.ModuleForTests("foo", "android_common")
 	manifest := m.Output("metalava.sbox.textproto")
-	sboxProto := android.RuleBuilderSboxProtoForTests(t, manifest)
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx, manifest)
 	manifestCommand := sboxProto.Commands[0].GetCommand()
 
 	// Api files are sorted from the narrowest api scope to the widest api scope.
@@ -2543,7 +2543,7 @@
 	`)
 	m := result.ModuleForTests(apiScopePublic.apiLibraryModuleName("foo"), "android_common")
 	manifest := m.Output("metalava.sbox.textproto")
-	sboxProto := android.RuleBuilderSboxProtoForTests(t, manifest)
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, result.TestContext, manifest)
 	manifestCommand := sboxProto.Commands[0].GetCommand()
 	classPathFlag := "--classpath __SBOX_SANDBOX_DIR__/out/.intermediates/bar/android_common/turbine-combined/bar.jar"
 	android.AssertStringDoesContain(t, "command expected to contain classpath flag", manifestCommand, classPathFlag)
diff --git a/java/lint_test.go b/java/lint_test.go
index ec901aa..5e6b8bc 100644
--- a/java/lint_test.go
+++ b/java/lint_test.go
@@ -39,7 +39,7 @@
 
 	foo := ctx.ModuleForTests("foo", "android_common")
 
-	sboxProto := android.RuleBuilderSboxProtoForTests(t, foo.Output("lint.sbox.textproto"))
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx, foo.Output("lint.sbox.textproto"))
 	if !strings.Contains(*sboxProto.Commands[0].Command, "--baseline lint-baseline.xml") {
 		t.Error("did not pass --baseline flag")
 	}
@@ -61,7 +61,7 @@
 
 	foo := ctx.ModuleForTests("foo", "android_common")
 
-	sboxProto := android.RuleBuilderSboxProtoForTests(t, foo.Output("lint.sbox.textproto"))
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx, foo.Output("lint.sbox.textproto"))
 	if strings.Contains(*sboxProto.Commands[0].Command, "--baseline") {
 		t.Error("passed --baseline flag for non existent file")
 	}
@@ -108,7 +108,7 @@
 
 	foo := ctx.ModuleForTests("foo", "android_common")
 
-	sboxProto := android.RuleBuilderSboxProtoForTests(t, foo.Output("lint.sbox.textproto"))
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, ctx, foo.Output("lint.sbox.textproto"))
 	if !strings.Contains(*sboxProto.Commands[0].Command, "--baseline mybaseline.xml") {
 		t.Error("did not use the correct file for baseline")
 	}
@@ -276,7 +276,7 @@
 			RunTestWithBp(t, thisBp)
 
 		foo := result.ModuleForTests("foo", "android_common")
-		sboxProto := android.RuleBuilderSboxProtoForTests(t, foo.Output("lint.sbox.textproto"))
+		sboxProto := android.RuleBuilderSboxProtoForTests(t, result.TestContext, foo.Output("lint.sbox.textproto"))
 		if !strings.Contains(*sboxProto.Commands[0].Command, "/"+testCase.expected_file) {
 			t.Error("did not use full api database for case", testCase)
 		}