diff --git a/rust/image_test.go b/rust/image_test.go
index 0581fa7..e5ecd79 100644
--- a/rust/image_test.go
+++ b/rust/image_test.go
@@ -45,7 +45,7 @@
 			}
 		`)
 
-	vendorBinary := ctx.ModuleForTests("fizz_vendor_available", "android_vendor_arm64_armv8-a").Module().(*cc.Module)
+	vendorBinary := ctx.ModuleForTests(t, "fizz_vendor_available", "android_vendor_arm64_armv8-a").Module().(*cc.Module)
 
 	if android.InList("libfoo_vendor.vendor", vendorBinary.Properties.AndroidMkStaticLibs) {
 		t.Errorf("vendorBinary should not have a staticlib dependency on libfoo_vendor.vendor: %#v", vendorBinary.Properties.AndroidMkStaticLibs)
@@ -64,7 +64,7 @@
 			}
 		`)
 
-	vendor := ctx.ModuleForTests("libfoo", "android_vendor_arm64_armv8-a_shared").Rule("rustc")
+	vendor := ctx.ModuleForTests(t, "libfoo", "android_vendor_arm64_armv8-a_shared").Rule("rustc")
 
 	if !strings.Contains(vendor.Args["rustcFlags"], "--cfg 'android_vndk'") {
 		t.Errorf("missing \"--cfg 'android_vndk'\" for libfoo vendor variant, rustcFlags: %#v", vendor.Args["rustcFlags"])
@@ -76,7 +76,7 @@
 		t.Errorf("unexpected \"--cfg 'android_product'\" for libfoo vendor variant, rustcFlags: %#v", vendor.Args["rustcFlags"])
 	}
 
-	product := ctx.ModuleForTests("libfoo", "android_product_arm64_armv8-a_shared").Rule("rustc")
+	product := ctx.ModuleForTests(t, "libfoo", "android_product_arm64_armv8-a_shared").Rule("rustc")
 	if !strings.Contains(product.Args["rustcFlags"], "--cfg 'android_vndk'") {
 		t.Errorf("missing \"--cfg 'android_vndk'\" for libfoo product variant, rustcFlags: %#v", product.Args["rustcFlags"])
 	}
@@ -87,7 +87,7 @@
 		t.Errorf("missing \"--cfg 'android_product'\" for libfoo product variant, rustcFlags: %#v", product.Args["rustcFlags"])
 	}
 
-	system := ctx.ModuleForTests("libfoo", "android_arm64_armv8-a_shared").Rule("rustc")
+	system := ctx.ModuleForTests(t, "libfoo", "android_arm64_armv8-a_shared").Rule("rustc")
 	if strings.Contains(system.Args["rustcFlags"], "--cfg 'android_vndk'") {
 		t.Errorf("unexpected \"--cfg 'android_vndk'\" for libfoo system variant, rustcFlags: %#v", system.Args["rustcFlags"])
 	}
@@ -119,7 +119,7 @@
 			}
 		`)
 
-	vendorRamdiskLibrary := ctx.ModuleForTests("libcc_vendor_ramdisk", "android_vendor_ramdisk_arm64_armv8-a_shared").Module().(*cc.Module)
+	vendorRamdiskLibrary := ctx.ModuleForTests(t, "libcc_vendor_ramdisk", "android_vendor_ramdisk_arm64_armv8-a_shared").Module().(*cc.Module)
 
 	if android.InList("libfoo_vendor_ramdisk.vendor_ramdisk", vendorRamdiskLibrary.Properties.AndroidMkStaticLibs) {
 		t.Errorf("libcc_vendor_ramdisk should not have a dependency on the libfoo_vendor_ramdisk static library")
@@ -144,7 +144,7 @@
 }
 
 func checkInstallPartition(t *testing.T, ctx *android.TestContext, name, variant, expected string) {
-	mod := ctx.ModuleForTests(name, variant).Module().(*Module)
+	mod := ctx.ModuleForTests(t, name, variant).Module().(*Module)
 	partitionDefined := false
 	checkPartition := func(specific bool, partition string) {
 		if specific {