Merge "Improve name of getModuleHeaderJarsAsNormalizedPaths"
diff --git a/android/androidmk.go b/android/androidmk.go
index 32d7712..9317567 100644
--- a/android/androidmk.go
+++ b/android/androidmk.go
@@ -44,6 +44,14 @@
 	ctx.RegisterSingletonType("androidmk", AndroidMkSingleton)
 }
 
+// Enable androidmk support.
+// * Register the singleton
+// * Configure that we are inside make
+var PrepareForTestWithAndroidMk = GroupFixturePreparers(
+	FixtureRegisterWithContext(RegisterAndroidMkBuildComponents),
+	FixtureModifyConfig(SetKatiEnabledForTests),
+)
+
 // Deprecated: Use AndroidMkEntriesProvider instead, especially if you're not going to use the
 // Custom function. It's easier to use and test.
 type AndroidMkDataProvider interface {
diff --git a/android/fixture.go b/android/fixture.go
index 25dee27..552c8f5 100644
--- a/android/fixture.go
+++ b/android/fixture.go
@@ -608,7 +608,7 @@
 	}
 }
 
-// AssertArrayString checks if the expected and actual values are equal using reflect.DeepEqual and
+// AssertDeepEquals checks if the expected and actual values are equal using reflect.DeepEqual and
 // if they are not then it reports an error prefixed with the supplied message and including a
 // reason for why it failed.
 func (h *TestHelper) AssertDeepEquals(message string, expected interface{}, actual interface{}) {
diff --git a/android/testing.go b/android/testing.go
index 03d7bd5..dd3d607 100644
--- a/android/testing.go
+++ b/android/testing.go
@@ -24,6 +24,7 @@
 	"testing"
 
 	"github.com/google/blueprint"
+	"github.com/google/blueprint/proptools"
 )
 
 func NewTestContext(config Config) *TestContext {
@@ -105,6 +106,17 @@
 	PrepareForTestWithAndroidBuildComponents,
 )
 
+// Prepares a test that may be missing dependencies by setting allow_missing_dependencies to
+// true.
+var PrepareForTestWithAllowMissingDependencies = GroupFixturePreparers(
+	FixtureModifyProductVariables(func(variables FixtureProductVariables) {
+		variables.Allow_missing_dependencies = proptools.BoolPtr(true)
+	}),
+	FixtureModifyContext(func(ctx *TestContext) {
+		ctx.SetAllowMissingDependencies(true)
+	}),
+)
+
 func NewTestArchContext(config Config) *TestContext {
 	ctx := NewTestContext(config)
 	ctx.preDeps = append(ctx.preDeps, registerArchMutator)
diff --git a/cc/cc_test.go b/cc/cc_test.go
index cc1f7d0..7d9fa47 100644
--- a/cc/cc_test.go
+++ b/cc/cc_test.go
@@ -55,7 +55,6 @@
 var ccFixtureFactory = android.NewFixtureFactory(
 	&buildDir,
 	PrepareForTestWithCcIncludeVndk,
-
 	android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
 		variables.DeviceVndkVersion = StringPtr("current")
 		variables.ProductVndkVersion = StringPtr("current")
@@ -3421,24 +3420,16 @@
 		}
 	`
 
-	config := TestConfig(buildDir, android.Android, nil, bp, nil)
-	config.TestProductVariables.Debuggable = BoolPtr(true)
+	result := ccFixtureFactory.Extend(
+		android.PrepareForTestWithVariables,
 
-	ctx := CreateTestContext(config)
-	ctx.PreDepsMutators(func(ctx android.RegisterMutatorsContext) {
-		ctx.BottomUp("variable", android.VariableMutator).Parallel()
-	})
-	ctx.Register()
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.Debuggable = BoolPtr(true)
+		}),
+	).RunTestWithBp(t, bp)
 
-	_, errs := ctx.ParseFileList(".", []string{"Android.bp"})
-	android.FailIfErrored(t, errs)
-	_, errs = ctx.PrepareBuildActions(config)
-	android.FailIfErrored(t, errs)
-
-	libfoo := ctx.ModuleForTests("libfoo", "android_arm64_armv8-a_static").Module().(*Module)
-	if !android.InList("-DBAR", libfoo.flags.Local.CppFlags) {
-		t.Errorf("expected -DBAR in cppflags, got %q", libfoo.flags.Local.CppFlags)
-	}
+	libfoo := result.Module("libfoo", "android_arm64_armv8-a_static").(*Module)
+	result.AssertStringListContains("cppflags", libfoo.flags.Local.CppFlags, "-DBAR")
 }
 
 func TestEmptyWholeStaticLibsAllowMissingDependencies(t *testing.T) {
@@ -3456,32 +3447,17 @@
 		}
 	`
 
-	config := TestConfig(buildDir, android.Android, nil, bp, nil)
-	config.TestProductVariables.Allow_missing_dependencies = BoolPtr(true)
+	result := ccFixtureFactory.Extend(
+		android.PrepareForTestWithAllowMissingDependencies,
+	).RunTestWithBp(t, bp)
 
-	ctx := CreateTestContext(config)
-	ctx.SetAllowMissingDependencies(true)
-	ctx.Register()
+	libbar := result.ModuleForTests("libbar", "android_arm64_armv8-a_static").Output("libbar.a")
+	result.AssertDeepEquals("libbar rule", android.ErrorRule, libbar.Rule)
 
-	_, errs := ctx.ParseFileList(".", []string{"Android.bp"})
-	android.FailIfErrored(t, errs)
-	_, errs = ctx.PrepareBuildActions(config)
-	android.FailIfErrored(t, errs)
+	result.AssertStringDoesContain("libbar error", libbar.Args["error"], "missing dependencies: libmissing")
 
-	libbar := ctx.ModuleForTests("libbar", "android_arm64_armv8-a_static").Output("libbar.a")
-	if g, w := libbar.Rule, android.ErrorRule; g != w {
-		t.Fatalf("Expected libbar rule to be %q, got %q", w, g)
-	}
-
-	if g, w := libbar.Args["error"], "missing dependencies: libmissing"; !strings.Contains(g, w) {
-		t.Errorf("Expected libbar error to contain %q, was %q", w, g)
-	}
-
-	libfoo := ctx.ModuleForTests("libfoo", "android_arm64_armv8-a_static").Output("libfoo.a")
-	if g, w := libfoo.Inputs.Strings(), libbar.Output.String(); !android.InList(w, g) {
-		t.Errorf("Expected libfoo.a to depend on %q, got %q", w, g)
-	}
-
+	libfoo := result.ModuleForTests("libfoo", "android_arm64_armv8-a_static").Output("libfoo.a")
+	result.AssertStringListContains("libfoo.a dependencies", libfoo.Inputs.Strings(), libbar.Output.String())
 }
 
 func TestInstallSharedLibs(t *testing.T) {
@@ -3671,8 +3647,9 @@
 	}
 }
 
-func makeMemtagTestConfig(t *testing.T) android.Config {
-	templateBp := `
+var prepareForTestWithMemtagHeap = android.GroupFixturePreparers(
+	android.FixtureModifyMockFS(func(fs android.MockFS) {
+		templateBp := `
 		cc_test {
 			name: "%[1]s_test",
 			gtest: false,
@@ -3726,35 +3703,30 @@
 			sanitize: { memtag_heap: true, diag: { memtag_heap: true }  },
 		}
 		`
-	subdirDefaultBp := fmt.Sprintf(templateBp, "default")
-	subdirExcludeBp := fmt.Sprintf(templateBp, "exclude")
-	subdirSyncBp := fmt.Sprintf(templateBp, "sync")
-	subdirAsyncBp := fmt.Sprintf(templateBp, "async")
+		subdirDefaultBp := fmt.Sprintf(templateBp, "default")
+		subdirExcludeBp := fmt.Sprintf(templateBp, "exclude")
+		subdirSyncBp := fmt.Sprintf(templateBp, "sync")
+		subdirAsyncBp := fmt.Sprintf(templateBp, "async")
 
-	mockFS := map[string][]byte{
-		"subdir_default/Android.bp": []byte(subdirDefaultBp),
-		"subdir_exclude/Android.bp": []byte(subdirExcludeBp),
-		"subdir_sync/Android.bp":    []byte(subdirSyncBp),
-		"subdir_async/Android.bp":   []byte(subdirAsyncBp),
-	}
-
-	return TestConfig(buildDir, android.Android, nil, "", mockFS)
-}
+		fs.Merge(android.MockFS{
+			"subdir_default/Android.bp": []byte(subdirDefaultBp),
+			"subdir_exclude/Android.bp": []byte(subdirExcludeBp),
+			"subdir_sync/Android.bp":    []byte(subdirSyncBp),
+			"subdir_async/Android.bp":   []byte(subdirAsyncBp),
+		})
+	}),
+	android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+		variables.MemtagHeapExcludePaths = []string{"subdir_exclude"}
+		variables.MemtagHeapSyncIncludePaths = []string{"subdir_sync"}
+		variables.MemtagHeapAsyncIncludePaths = []string{"subdir_async"}
+	}),
+)
 
 func TestSanitizeMemtagHeap(t *testing.T) {
 	variant := "android_arm64_armv8-a"
 
-	config := makeMemtagTestConfig(t)
-	config.TestProductVariables.MemtagHeapExcludePaths = []string{"subdir_exclude"}
-	config.TestProductVariables.MemtagHeapSyncIncludePaths = []string{"subdir_sync"}
-	config.TestProductVariables.MemtagHeapAsyncIncludePaths = []string{"subdir_async"}
-	ctx := CreateTestContext(config)
-	ctx.Register()
-
-	_, errs := ctx.ParseFileList(".", []string{"Android.bp", "subdir_default/Android.bp", "subdir_exclude/Android.bp", "subdir_sync/Android.bp", "subdir_async/Android.bp"})
-	android.FailIfErrored(t, errs)
-	_, errs = ctx.PrepareBuildActions(config)
-	android.FailIfErrored(t, errs)
+	result := ccFixtureFactory.Extend(prepareForTestWithMemtagHeap).RunTest(t)
+	ctx := result.TestContext
 
 	checkHasMemtagNote(t, ctx.ModuleForTests("default_test", variant), Sync)
 	checkHasMemtagNote(t, ctx.ModuleForTests("default_test_false", variant), None)
@@ -3808,18 +3780,13 @@
 func TestSanitizeMemtagHeapWithSanitizeDevice(t *testing.T) {
 	variant := "android_arm64_armv8-a"
 
-	config := makeMemtagTestConfig(t)
-	config.TestProductVariables.MemtagHeapExcludePaths = []string{"subdir_exclude"}
-	config.TestProductVariables.MemtagHeapSyncIncludePaths = []string{"subdir_sync"}
-	config.TestProductVariables.MemtagHeapAsyncIncludePaths = []string{"subdir_async"}
-	config.TestProductVariables.SanitizeDevice = []string{"memtag_heap"}
-	ctx := CreateTestContext(config)
-	ctx.Register()
-
-	_, errs := ctx.ParseFileList(".", []string{"Android.bp", "subdir_default/Android.bp", "subdir_exclude/Android.bp", "subdir_sync/Android.bp", "subdir_async/Android.bp"})
-	android.FailIfErrored(t, errs)
-	_, errs = ctx.PrepareBuildActions(config)
-	android.FailIfErrored(t, errs)
+	result := ccFixtureFactory.Extend(
+		prepareForTestWithMemtagHeap,
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.SanitizeDevice = []string{"memtag_heap"}
+		}),
+	).RunTest(t)
+	ctx := result.TestContext
 
 	checkHasMemtagNote(t, ctx.ModuleForTests("default_test", variant), Sync)
 	checkHasMemtagNote(t, ctx.ModuleForTests("default_test_false", variant), None)
@@ -3873,19 +3840,14 @@
 func TestSanitizeMemtagHeapWithSanitizeDeviceDiag(t *testing.T) {
 	variant := "android_arm64_armv8-a"
 
-	config := makeMemtagTestConfig(t)
-	config.TestProductVariables.MemtagHeapExcludePaths = []string{"subdir_exclude"}
-	config.TestProductVariables.MemtagHeapSyncIncludePaths = []string{"subdir_sync"}
-	config.TestProductVariables.MemtagHeapAsyncIncludePaths = []string{"subdir_async"}
-	config.TestProductVariables.SanitizeDevice = []string{"memtag_heap"}
-	config.TestProductVariables.SanitizeDeviceDiag = []string{"memtag_heap"}
-	ctx := CreateTestContext(config)
-	ctx.Register()
-
-	_, errs := ctx.ParseFileList(".", []string{"Android.bp", "subdir_default/Android.bp", "subdir_exclude/Android.bp", "subdir_sync/Android.bp", "subdir_async/Android.bp"})
-	android.FailIfErrored(t, errs)
-	_, errs = ctx.PrepareBuildActions(config)
-	android.FailIfErrored(t, errs)
+	result := ccFixtureFactory.Extend(
+		prepareForTestWithMemtagHeap,
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.SanitizeDevice = []string{"memtag_heap"}
+			variables.SanitizeDeviceDiag = []string{"memtag_heap"}
+		}),
+	).RunTest(t)
+	ctx := result.TestContext
 
 	checkHasMemtagNote(t, ctx.ModuleForTests("default_test", variant), Sync)
 	checkHasMemtagNote(t, ctx.ModuleForTests("default_test_false", variant), None)
diff --git a/cc/prebuilt_test.go b/cc/prebuilt_test.go
index ee4de6e..20274b2 100644
--- a/cc/prebuilt_test.go
+++ b/cc/prebuilt_test.go
@@ -23,27 +23,16 @@
 	"github.com/google/blueprint"
 )
 
-func testPrebuilt(t *testing.T, bp string, fs map[string][]byte, handlers ...configCustomizer) *android.TestContext {
-	config := TestConfig(buildDir, android.Android, nil, bp, fs)
-	ctx := CreateTestContext(config)
+var prebuiltFixtureFactory = ccFixtureFactory.Extend(
+	android.PrepareForTestWithAndroidMk,
+)
 
-	// Enable androidmk support.
-	// * Register the singleton
-	// * Configure that we are inside make
-	// * Add CommonOS to ensure that androidmk processing works.
-	android.RegisterAndroidMkBuildComponents(ctx)
-	android.SetKatiEnabledForTests(config)
+func testPrebuilt(t *testing.T, bp string, fs android.MockFS, handlers ...android.FixturePreparer) *android.TestContext {
+	result := prebuiltFixtureFactory.Extend(
+		fs.AddToFixture(),
+	).Extend(handlers...).RunTestWithBp(t, bp)
 
-	for _, handler := range handlers {
-		handler(config)
-	}
-
-	ctx.Register()
-	_, errs := ctx.ParseFileList(".", []string{"Android.bp"})
-	android.FailIfErrored(t, errs)
-	_, errs = ctx.PrepareBuildActions(config)
-	android.FailIfErrored(t, errs)
-	return ctx
+	return result.TestContext
 }
 
 type configCustomizer func(config android.Config)
@@ -370,9 +359,11 @@
 	assertString(t, static2.OutputFile().Path().Base(), "libf.a")
 
 	// With SANITIZE_TARGET=hwaddress
-	ctx = testPrebuilt(t, bp, fs, func(config android.Config) {
-		config.TestProductVariables.SanitizeDevice = []string{"hwaddress"}
-	})
+	ctx = testPrebuilt(t, bp, fs,
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.SanitizeDevice = []string{"hwaddress"}
+		}),
+	)
 
 	shared_rule = ctx.ModuleForTests("libtest", "android_arm64_armv8-a_shared_hwasan").Rule("android/soong/cc.strip")
 	assertString(t, shared_rule.Input.String(), "hwasan/libf.so")
diff --git a/filesystem/bootimg.go b/filesystem/bootimg.go
index 764f045..8e2c554 100644
--- a/filesystem/bootimg.go
+++ b/filesystem/bootimg.go
@@ -237,3 +237,13 @@
 func (b *bootimg) OutputPath() android.Path {
 	return b.output
 }
+
+var _ android.OutputFileProducer = (*bootimg)(nil)
+
+// Implements android.OutputFileProducer
+func (b *bootimg) OutputFiles(tag string) (android.Paths, error) {
+	if tag == "" {
+		return []android.Path{b.output}, nil
+	}
+	return nil, fmt.Errorf("unsupported module reference tag %q", tag)
+}
diff --git a/filesystem/logical_partition.go b/filesystem/logical_partition.go
index e547203..16b6037 100644
--- a/filesystem/logical_partition.go
+++ b/filesystem/logical_partition.go
@@ -208,3 +208,13 @@
 func (l *logicalPartition) OutputPath() android.Path {
 	return l.output
 }
+
+var _ android.OutputFileProducer = (*logicalPartition)(nil)
+
+// Implements android.OutputFileProducer
+func (l *logicalPartition) OutputFiles(tag string) (android.Paths, error) {
+	if tag == "" {
+		return []android.Path{l.output}, nil
+	}
+	return nil, fmt.Errorf("unsupported module reference tag %q", tag)
+}
diff --git a/java/lint.go b/java/lint.go
index 5940eac..9f677db 100644
--- a/java/lint.go
+++ b/java/lint.go
@@ -189,6 +189,10 @@
 	remoteRSPInputs android.Paths
 }
 
+func lintRBEExecStrategy(ctx android.ModuleContext) string {
+	return ctx.Config().GetenvWithDefault("RBE_LINT_EXEC_STRATEGY", remoteexec.LocalExecStrategy)
+}
+
 func (l *linter) writeLintProjectXML(ctx android.ModuleContext, rule *android.RuleBuilder) lintPaths {
 	var deps android.Paths
 	var remoteInputs android.Paths
@@ -280,7 +284,8 @@
 	cmd.FlagForEachArg("--extra_checks_jar ", l.extraLintCheckJars.Strings())
 	trackInputDependency(l.extraLintCheckJars...)
 
-	if ctx.Config().UseRBE() && ctx.Config().IsEnvTrue("RBE_LINT") {
+	if ctx.Config().UseRBE() && ctx.Config().IsEnvTrue("RBE_LINT") &&
+		lintRBEExecStrategy(ctx) != remoteexec.LocalExecStrategy {
 		// TODO(b/181912787): remove these and use "." instead.
 		cmd.FlagWithArg("--root_dir ", "/b/f/w")
 	} else {
@@ -391,7 +396,7 @@
 		pool := ctx.Config().GetenvWithDefault("RBE_LINT_POOL", "java16")
 		// TODO(b/181912787): this should be local fallback once the hack that passes /b/f/w in project.xml
 		// is removed.
-		execStrategy := ctx.Config().GetenvWithDefault("RBE_LINT_EXEC_STRATEGY", remoteexec.RemoteExecStrategy)
+		execStrategy := lintRBEExecStrategy(ctx)
 		labels := map[string]string{"type": "tool", "name": "lint"}
 		rule.Remoteable(android.RemoteRuleSupports{RBE: true})
 		remoteInputs := lintPaths.remoteInputs
@@ -417,6 +422,7 @@
 				"ANDROID_SDK_HOME",
 				"SDK_ANNOTATIONS",
 				"LINT_OPTS",
+				"LANG",
 			},
 			Platform: map[string]string{remoteexec.PoolKey: pool},
 		}).NoVarTemplate(ctx.Config()))
diff --git a/rust/testing.go b/rust/testing.go
index 9534ab5..5be71c9 100644
--- a/rust/testing.go
+++ b/rust/testing.go
@@ -20,6 +20,30 @@
 	"android/soong/genrule"
 )
 
+// Preparer that will define all cc module types and a limited set of mutators and singletons that
+// make those module types usable.
+var PrepareForTestWithRustBuildComponents = android.GroupFixturePreparers(
+	android.FixtureRegisterWithContext(RegisterRequiredBuildComponentsForTest),
+)
+
+// The directory in which rust test default modules will be defined.
+//
+// Placing them here ensures that their location does not conflict with default test modules
+// defined by other packages.
+const rustDefaultsDir = "defaults/rust/"
+
+// Preparer that will define default rust modules, e.g. standard prebuilt modules.
+var PrepareForTestWithRustDefaultModules = android.GroupFixturePreparers(
+	cc.PrepareForTestWithCcDefaultModules,
+	PrepareForTestWithRustBuildComponents,
+	android.FixtureAddTextFile(rustDefaultsDir+"Android.bp", GatherRequiredDepsForTest()),
+)
+
+// Preparer that will allow use of all rust modules fully.
+var PrepareForIntegrationTestWithRust = android.GroupFixturePreparers(
+	PrepareForTestWithRustDefaultModules,
+)
+
 func GatherRequiredDepsForTest() string {
 	bp := `
 		rust_prebuilt_library {
diff --git a/sdk/cc_sdk_test.go b/sdk/cc_sdk_test.go
index 3591777..6da135a 100644
--- a/sdk/cc_sdk_test.go
+++ b/sdk/cc_sdk_test.go
@@ -101,7 +101,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -353,7 +353,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -440,7 +440,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAllCopyRules(`
 myinclude/Test.h -> include/myinclude/Test.h
 .intermediates/mynativelib1/android_arm64_armv8-a_shared/mynativelib1.so -> arm64/lib/mynativelib1.so
@@ -486,7 +486,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -556,7 +556,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -615,7 +615,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mymodule_exports", "",
+	CheckSnapshot(result, "mymodule_exports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -700,7 +700,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -853,7 +853,7 @@
 
 	result := runTests(t, ctx, config)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -991,7 +991,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mymodule_exports", "",
+	CheckSnapshot(result, "mymodule_exports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1099,7 +1099,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1200,7 +1200,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1297,7 +1297,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1424,7 +1424,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1552,7 +1552,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1615,7 +1615,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1729,7 +1729,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1843,7 +1843,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1940,7 +1940,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1978,7 +1978,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -2080,7 +2080,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -2193,7 +2193,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -2266,7 +2266,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -2377,7 +2377,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -2430,7 +2430,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -2543,7 +2543,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -2658,7 +2658,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkUnversionedAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
diff --git a/sdk/exports_test.go b/sdk/exports_test.go
index 1c59244..54a40d2 100644
--- a/sdk/exports_test.go
+++ b/sdk/exports_test.go
@@ -42,7 +42,7 @@
 			"package/Android.bp": []byte(packageBp),
 		})
 
-	result.CheckSnapshot("myexports", "package",
+	CheckSnapshot(result, "myexports", "package",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
diff --git a/sdk/java_sdk_test.go b/sdk/java_sdk_test.go
index 111b22c..ef8e4a0 100644
--- a/sdk/java_sdk_test.go
+++ b/sdk/java_sdk_test.go
@@ -125,9 +125,9 @@
 	`)
 
 	// Make sure that the mysdk module depends on "sdkmember" and not "prebuilt_sdkmember".
-	java.CheckModuleDependencies(t, result.ctx, "mysdk", "android_common", []string{"sdkmember"})
+	java.CheckModuleDependencies(t, result.TestContext, "mysdk", "android_common", []string{"sdkmember"})
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`// This is auto-generated. DO NOT EDIT.
 
 java_import {
@@ -224,11 +224,11 @@
 		}
 	`)
 
-	sdkMemberV1 := result.ctx.ModuleForTests("sdkmember_mysdk_1", "android_common").Rule("combineJar").Output
-	sdkMemberV2 := result.ctx.ModuleForTests("sdkmember_mysdk_2", "android_common").Rule("combineJar").Output
+	sdkMemberV1 := result.ModuleForTests("sdkmember_mysdk_1", "android_common").Rule("combineJar").Output
+	sdkMemberV2 := result.ModuleForTests("sdkmember_mysdk_2", "android_common").Rule("combineJar").Output
 
-	javalibForMyApex := result.ctx.ModuleForTests("myjavalib", "android_common_apex10000_mysdk_1")
-	javalibForMyApex2 := result.ctx.ModuleForTests("myjavalib", "android_common_apex10000_mysdk_2")
+	javalibForMyApex := result.ModuleForTests("myjavalib", "android_common_apex10000_mysdk_1")
+	javalibForMyApex2 := result.ModuleForTests("myjavalib", "android_common_apex10000_mysdk_2")
 
 	// Depending on the uses_sdks value, different libs are linked
 	ensureListContains(t, pathsToStrings(javalibForMyApex.Rule("javac").Implicits), sdkMemberV1.String())
@@ -255,7 +255,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -312,7 +312,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -369,7 +369,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -440,7 +440,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -496,7 +496,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -551,7 +551,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -607,7 +607,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -662,7 +662,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -731,7 +731,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -827,7 +827,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -918,7 +918,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("myexports", "",
+	CheckSnapshot(result, "myexports", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1032,7 +1032,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1133,7 +1133,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1202,7 +1202,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1274,7 +1274,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1367,7 +1367,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1475,7 +1475,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1563,7 +1563,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -1639,7 +1639,7 @@
 		}
 	`)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
diff --git a/sdk/sdk_test.go b/sdk/sdk_test.go
index a7acd0c..65a9001 100644
--- a/sdk/sdk_test.go
+++ b/sdk/sdk_test.go
@@ -169,7 +169,7 @@
 			"package/Android.bp": []byte(packageBp),
 		})
 
-	result.CheckSnapshot("mysdk", "package",
+	CheckSnapshot(result, "mysdk", "package",
 		checkAndroidBpContents(`
 // This is auto-generated. DO NOT EDIT.
 
@@ -317,7 +317,7 @@
 	`
 	result := testSdkWithFs(t, sdk, nil)
 
-	result.CheckSnapshot("mysdk", "",
+	CheckSnapshot(result, "mysdk", "",
 		checkAllOtherCopyRules(`.intermediates/mysdk/common_os/mysdk-current.zip -> mysdk-current.zip`),
 	)
 }
diff --git a/sdk/testing.go b/sdk/testing.go
index e291bdb..3fb27ca 100644
--- a/sdk/testing.go
+++ b/sdk/testing.go
@@ -136,9 +136,8 @@
 	_, errs = ctx.PrepareBuildActions(config)
 	android.FailIfErrored(t, errs)
 	return &testSdkResult{
-		TestHelper: android.TestHelper{T: t},
-		ctx:        ctx,
-		config:     config,
+		TestHelper:  android.TestHelper{T: t},
+		TestContext: ctx,
 	}
 }
 
@@ -184,17 +183,20 @@
 // checking the state of the build structures.
 type testSdkResult struct {
 	android.TestHelper
-	ctx    *android.TestContext
-	config android.Config
+	*android.TestContext
+}
+
+func (result *testSdkResult) Module(name string, variant string) android.Module {
+	return result.ModuleForTests(name, variant).Module()
 }
 
 // Analyse the sdk build rules to extract information about what it is doing.
-
+//
 // e.g. find the src/dest pairs from each cp command, the various zip files
 // generated, etc.
-func (r *testSdkResult) getSdkSnapshotBuildInfo(sdk *sdk) *snapshotBuildInfo {
+func getSdkSnapshotBuildInfo(result *testSdkResult, sdk *sdk) *snapshotBuildInfo {
 	info := &snapshotBuildInfo{
-		r:                            r,
+		r:                            result,
 		androidBpContents:            sdk.GetAndroidBpContentsForTests(),
 		androidUnversionedBpContents: sdk.GetUnversionedAndroidBpContentsForTests(),
 		androidVersionedBpContents:   sdk.GetVersionedAndroidBpContentsForTests(),
@@ -237,7 +239,7 @@
 			info.intermediateZip = info.outputZip
 			mergeInput := android.NormalizePathForTesting(bp.Input)
 			if info.intermediateZip != mergeInput {
-				r.Errorf("Expected intermediate zip %s to be an input to merge zips but found %s instead",
+				result.Errorf("Expected intermediate zip %s to be an input to merge zips but found %s instead",
 					info.intermediateZip, mergeInput)
 			}
 
@@ -256,28 +258,20 @@
 	return info
 }
 
-func (r *testSdkResult) Module(name string, variant string) android.Module {
-	return r.ctx.ModuleForTests(name, variant).Module()
-}
-
-func (r *testSdkResult) ModuleForTests(name string, variant string) android.TestingModule {
-	return r.ctx.ModuleForTests(name, variant)
-}
-
 // Check the snapshot build rules.
 //
 // Takes a list of functions which check different facets of the snapshot build rules.
 // Allows each test to customize what is checked without duplicating lots of code
 // or proliferating check methods of different flavors.
-func (r *testSdkResult) CheckSnapshot(name string, dir string, checkers ...snapshotBuildInfoChecker) {
-	r.Helper()
+func CheckSnapshot(result *testSdkResult, name string, dir string, checkers ...snapshotBuildInfoChecker) {
+	result.Helper()
 
 	// The sdk CommonOS variant is always responsible for generating the snapshot.
 	variant := android.CommonOS.Name
 
-	sdk := r.Module(name, variant).(*sdk)
+	sdk := result.Module(name, variant).(*sdk)
 
-	snapshotBuildInfo := r.getSdkSnapshotBuildInfo(sdk)
+	snapshotBuildInfo := getSdkSnapshotBuildInfo(result, sdk)
 
 	// Check state of the snapshot build.
 	for _, checker := range checkers {
@@ -289,7 +283,7 @@
 	if dir != "" {
 		dir = filepath.Clean(dir) + "/"
 	}
-	r.AssertStringEquals("Snapshot zip file in wrong place",
+	result.AssertStringEquals("Snapshot zip file in wrong place",
 		fmt.Sprintf(".intermediates/%s%s/%s/%s-current.zip", dir, name, variant, name), actual)
 
 	// Populate a mock filesystem with the files that would have been copied by
@@ -300,7 +294,7 @@
 	}
 
 	// Process the generated bp file to make sure it is valid.
-	testSdkWithFs(r.T, snapshotBuildInfo.androidBpContents, fs)
+	testSdkWithFs(result.T, snapshotBuildInfo.androidBpContents, fs)
 }
 
 type snapshotBuildInfoChecker func(info *snapshotBuildInfo)