diff --git a/tradefed/autogen_bazel.go b/tradefed/autogen_bazel.go
index 8283984..cc16176 100644
--- a/tradefed/autogen_bazel.go
+++ b/tradefed/autogen_bazel.go
@@ -49,7 +49,7 @@
 }
 
 func GetTestConfigAttributes(
-	ctx android.TopDownMutatorContext,
+	ctx android.Bp2buildMutatorContext,
 	testConfig *string,
 	extraTestConfigs []string,
 	autoGenConfig *bool,
@@ -93,7 +93,7 @@
 }
 
 func GetTestConfig(
-	ctx android.TopDownMutatorContext,
+	ctx android.Bp2buildMutatorContext,
 	testConfig *string,
 ) *bazel.Label {