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")
 }