diff --git a/java/java_test.go b/java/java_test.go
index 4a549e6..49838c7 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -56,11 +56,11 @@
 	os.Exit(run())
 }
 
-func testConfig(env map[string]string) android.Config {
-	return TestConfig(buildDir, env)
+func testConfig(env map[string]string, bp string, fs map[string][]byte) android.Config {
+	return TestConfig(buildDir, env, bp, fs)
 }
 
-func testContext(bp string, fs map[string][]byte) *android.TestContext {
+func testContext() *android.TestContext {
 
 	ctx := android.NewTestArchContext()
 	ctx.RegisterModuleType("android_app", AndroidAppFactory)
@@ -116,119 +116,16 @@
 		ctx.BottomUp("begin", cc.BeginMutator).Parallel()
 	})
 
-	bp += GatherRequiredDepsForTest()
-
-	mockFS := map[string][]byte{
-		"Android.bp":             []byte(bp),
-		"a.java":                 nil,
-		"b.java":                 nil,
-		"c.java":                 nil,
-		"b.kt":                   nil,
-		"a.jar":                  nil,
-		"b.jar":                  nil,
-		"APP_NOTICE":             nil,
-		"GENRULE_NOTICE":         nil,
-		"LIB_NOTICE":             nil,
-		"TOOL_NOTICE":            nil,
-		"java-res/a/a":           nil,
-		"java-res/b/b":           nil,
-		"java-res2/a":            nil,
-		"java-fg/a.java":         nil,
-		"java-fg/b.java":         nil,
-		"java-fg/c.java":         nil,
-		"api/current.txt":        nil,
-		"api/removed.txt":        nil,
-		"api/system-current.txt": nil,
-		"api/system-removed.txt": nil,
-		"api/test-current.txt":   nil,
-		"api/test-removed.txt":   nil,
-		"framework/aidl/a.aidl":  nil,
-
-		"prebuilts/ndk/current/sources/cxx-stl/llvm-libc++/libs/arm64-v8a/libc++_shared.so": nil,
-
-		"prebuilts/sdk/14/public/android.jar":         nil,
-		"prebuilts/sdk/14/public/framework.aidl":      nil,
-		"prebuilts/sdk/14/system/android.jar":         nil,
-		"prebuilts/sdk/17/public/android.jar":         nil,
-		"prebuilts/sdk/17/public/framework.aidl":      nil,
-		"prebuilts/sdk/17/system/android.jar":         nil,
-		"prebuilts/sdk/29/public/android.jar":         nil,
-		"prebuilts/sdk/29/public/framework.aidl":      nil,
-		"prebuilts/sdk/29/system/android.jar":         nil,
-		"prebuilts/sdk/current/core/android.jar":      nil,
-		"prebuilts/sdk/current/public/android.jar":    nil,
-		"prebuilts/sdk/current/public/framework.aidl": nil,
-		"prebuilts/sdk/current/public/core.jar":       nil,
-		"prebuilts/sdk/current/system/android.jar":    nil,
-		"prebuilts/sdk/current/test/android.jar":      nil,
-		"prebuilts/sdk/28/public/api/foo.txt":         nil,
-		"prebuilts/sdk/28/system/api/foo.txt":         nil,
-		"prebuilts/sdk/28/test/api/foo.txt":           nil,
-		"prebuilts/sdk/28/public/api/foo-removed.txt": nil,
-		"prebuilts/sdk/28/system/api/foo-removed.txt": nil,
-		"prebuilts/sdk/28/test/api/foo-removed.txt":   nil,
-		"prebuilts/sdk/28/public/api/bar.txt":         nil,
-		"prebuilts/sdk/28/system/api/bar.txt":         nil,
-		"prebuilts/sdk/28/test/api/bar.txt":           nil,
-		"prebuilts/sdk/28/public/api/bar-removed.txt": nil,
-		"prebuilts/sdk/28/system/api/bar-removed.txt": nil,
-		"prebuilts/sdk/28/test/api/bar-removed.txt":   nil,
-		"prebuilts/sdk/tools/core-lambda-stubs.jar":   nil,
-		"prebuilts/sdk/Android.bp":                    []byte(`prebuilt_apis { name: "sdk", api_dirs: ["14", "28", "current"],}`),
-
-		"prebuilts/apk/app.apk":        nil,
-		"prebuilts/apk/app_arm.apk":    nil,
-		"prebuilts/apk/app_arm64.apk":  nil,
-		"prebuilts/apk/app_xhdpi.apk":  nil,
-		"prebuilts/apk/app_xxhdpi.apk": nil,
-
-		// For framework-res, which is an implicit dependency for framework
-		"AndroidManifest.xml":                        nil,
-		"build/make/target/product/security/testkey": nil,
-
-		"build/soong/scripts/jar-wrapper.sh": nil,
-
-		"build/make/core/verify_uses_libraries.sh": nil,
-
-		"build/make/core/proguard.flags":             nil,
-		"build/make/core/proguard_basic_keeps.flags": nil,
-
-		"jdk8/jre/lib/jce.jar": nil,
-		"jdk8/jre/lib/rt.jar":  nil,
-		"jdk8/lib/tools.jar":   nil,
-
-		"bar-doc/a.java":                 nil,
-		"bar-doc/b.java":                 nil,
-		"bar-doc/IFoo.aidl":              nil,
-		"bar-doc/IBar.aidl":              nil,
-		"bar-doc/known_oj_tags.txt":      nil,
-		"external/doclava/templates-sdk": nil,
-
-		"cert/new_cert.x509.pem": nil,
-		"cert/new_cert.pk8":      nil,
-
-		"testdata/data": nil,
-
-		"stubs-sources/foo/Foo.java": nil,
-		"stubs/sources/foo/Foo.java": nil,
-	}
-
-	for k, v := range fs {
-		mockFS[k] = v
-	}
-
-	ctx.MockFileSystem(mockFS)
-
 	return ctx
 }
 
 func run(t *testing.T, ctx *android.TestContext, config android.Config) {
 	t.Helper()
 
-	pathCtx := android.PathContextForTesting(config, nil)
+	pathCtx := android.PathContextForTesting(config)
 	setDexpreoptTestGlobalConfig(config, dexpreopt.GlobalConfigForTests(pathCtx))
 
-	ctx.Register()
+	ctx.Register(config)
 	_, errs := ctx.ParseBlueprintsFiles("Android.bp")
 	android.FailIfErrored(t, errs)
 	_, errs = ctx.PrepareBuildActions(config)
@@ -237,17 +134,17 @@
 
 func testJavaError(t *testing.T, pattern string, bp string) (*android.TestContext, android.Config) {
 	t.Helper()
-	return testJavaErrorWithConfig(t, pattern, bp, testConfig(nil))
+	return testJavaErrorWithConfig(t, pattern, testConfig(nil, bp, nil))
 }
 
-func testJavaErrorWithConfig(t *testing.T, pattern string, bp string, config android.Config) (*android.TestContext, android.Config) {
+func testJavaErrorWithConfig(t *testing.T, pattern string, config android.Config) (*android.TestContext, android.Config) {
 	t.Helper()
-	ctx := testContext(bp, nil)
+	ctx := testContext()
 
-	pathCtx := android.PathContextForTesting(config, nil)
+	pathCtx := android.PathContextForTesting(config)
 	setDexpreoptTestGlobalConfig(config, dexpreopt.GlobalConfigForTests(pathCtx))
 
-	ctx.Register()
+	ctx.Register(config)
 	_, errs := ctx.ParseBlueprintsFiles("Android.bp")
 	if len(errs) > 0 {
 		android.FailIfNoMatchingErrors(t, pattern, errs)
@@ -266,12 +163,12 @@
 
 func testJava(t *testing.T, bp string) (*android.TestContext, android.Config) {
 	t.Helper()
-	return testJavaWithConfig(t, bp, testConfig(nil))
+	return testJavaWithConfig(t, testConfig(nil, bp, nil))
 }
 
-func testJavaWithConfig(t *testing.T, bp string, config android.Config) (*android.TestContext, android.Config) {
+func testJavaWithConfig(t *testing.T, config android.Config) (*android.TestContext, android.Config) {
 	t.Helper()
-	ctx := testContext(bp, nil)
+	ctx := testContext()
 	run(t, ctx, config)
 
 	return ctx, config
@@ -520,9 +417,6 @@
 	`)
 
 	for _, enforce := range []bool{true, false} {
-
-		config := testConfig(nil)
-		config.TestProductVariables.EnforceProductPartitionInterface = proptools.BoolPtr(enforce)
 		bp := `
 			java_library {
 				name: "foo",
@@ -530,10 +424,13 @@
 				product_specific: true,
 			}
 		`
+
+		config := testConfig(nil, bp, nil)
+		config.TestProductVariables.EnforceProductPartitionInterface = proptools.BoolPtr(enforce)
 		if enforce {
-			testJavaErrorWithConfig(t, "sdk_version must have a value when the module is located at vendor or product", bp, config)
+			testJavaErrorWithConfig(t, "sdk_version must have a value when the module is located at vendor or product", config)
 		} else {
-			testJavaWithConfig(t, bp, config)
+			testJavaWithConfig(t, config)
 		}
 	}
 }
@@ -1127,8 +1024,7 @@
 }
 
 func TestJavaLibrary(t *testing.T) {
-	config := testConfig(nil)
-	ctx := testContext("", map[string][]byte{
+	config := testConfig(nil, "", map[string][]byte{
 		"libcore/Android.bp": []byte(`
 				java_library {
 						name: "core",
@@ -1136,6 +1032,7 @@
 						system_modules: "none",
 				}`),
 	})
+	ctx := testContext()
 	run(t, ctx, config)
 }