diff --git a/python/python_test.go b/python/python_test.go
index 57aaa34..4c30d95 100644
--- a/python/python_test.go
+++ b/python/python_test.go
@@ -26,8 +26,6 @@
 	"testing"
 
 	"android/soong/android"
-
-	"github.com/google/blueprint"
 )
 
 type pyBinary struct {
@@ -306,17 +304,17 @@
 func TestPythonModule(t *testing.T) {
 	config, buildDir := setupBuildEnv(t)
 	defer tearDownBuildEnv(buildDir)
-	android.TestPreDepsMutators(func(ctx android.RegisterMutatorsContext) {
-		ctx.BottomUp("version_split", versionSplitMutator()).Parallel()
-	})
 	for _, d := range data {
 		t.Run(d.desc, func(t *testing.T) {
-			ctx := blueprint.NewContext()
-			android.RegisterTestMutators(ctx)
+			ctx := android.NewTestContext()
+			ctx.PreDepsMutators(func(ctx android.RegisterMutatorsContext) {
+				ctx.BottomUp("version_split", versionSplitMutator()).Parallel()
+			})
 			ctx.RegisterModuleType("python_library_host",
 				android.ModuleFactoryAdaptor(PythonLibraryHostFactory))
 			ctx.RegisterModuleType("python_binary_host",
 				android.ModuleFactoryAdaptor(PythonBinaryHostFactory))
+			ctx.Register()
 			ctx.MockFileSystem(d.mockFiles)
 			_, testErrs := ctx.ParseBlueprintsFiles(bpFile)
 			fail(t, testErrs)
@@ -360,15 +358,12 @@
 	return
 }
 
-func expectModule(t *testing.T, ctx *blueprint.Context, buildDir, name, variant string,
+func expectModule(t *testing.T, ctx *android.TestContext, buildDir, name, variant string,
 	expPyRunfiles, expDepsPyRunfiles []string,
 	expParSpec string, expDepsParSpecs []string) (testErrs []error) {
-	module := findModule(ctx, name, variant)
-	if module == nil {
-		t.Fatalf("failed to find module %s!", name)
-	}
+	module := ctx.ModuleForTests(name, variant)
 
-	base, baseOk := module.(*pythonBaseModule)
+	base, baseOk := module.Module().(*pythonBaseModule)
 	if !baseOk {
 		t.Fatalf("%s is not Python module!", name)
 	}
@@ -438,16 +433,6 @@
 	os.RemoveAll(buildDir)
 }
 
-func findModule(ctx *blueprint.Context, name, variant string) blueprint.Module {
-	var ret blueprint.Module
-	ctx.VisitAllModules(func(m blueprint.Module) {
-		if ctx.ModuleName(m) == name && ctx.ModuleSubDir(m) == variant {
-			ret = m
-		}
-	})
-	return ret
-}
-
 func fail(t *testing.T, errs []error) {
 	if len(errs) > 0 {
 		for _, err := range errs {