diff --git a/java/java_test.go b/java/java_test.go
index b68945f..75813e0 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -25,12 +25,12 @@
 	"strings"
 	"testing"
 
-	"android/soong/genrule"
 	"github.com/google/blueprint/proptools"
 
 	"android/soong/android"
 	"android/soong/cc"
 	"android/soong/dexpreopt"
+	"android/soong/genrule"
 	"android/soong/python"
 )
 
@@ -1296,9 +1296,9 @@
 	})
 
 	foo := ctx.ModuleForTests("foo", "android_common")
-	rule := foo.Rule("lint")
 
-	if !strings.Contains(rule.RuleParams.Command, "--baseline lint-baseline.xml") {
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, foo.Output("lint.sbox.textproto"))
+	if !strings.Contains(*sboxProto.Commands[0].Command, "--baseline lint-baseline.xml") {
 		t.Error("did not pass --baseline flag")
 	}
 }
@@ -1318,9 +1318,9 @@
        `, map[string][]byte{})
 
 	foo := ctx.ModuleForTests("foo", "android_common")
-	rule := foo.Rule("lint")
 
-	if strings.Contains(rule.RuleParams.Command, "--baseline") {
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, foo.Output("lint.sbox.textproto"))
+	if strings.Contains(*sboxProto.Commands[0].Command, "--baseline") {
 		t.Error("passed --baseline flag for non existent file")
 	}
 }
@@ -1376,9 +1376,9 @@
 	})
 
 	foo := ctx.ModuleForTests("foo", "android_common")
-	rule := foo.Rule("lint")
 
-	if !strings.Contains(rule.RuleParams.Command, "--baseline mybaseline.xml") {
+	sboxProto := android.RuleBuilderSboxProtoForTests(t, 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")
 	}
 }