diff --git a/Android.bp b/Android.bp
index 496dcd286..250c969e4 100644
--- a/Android.bp
+++ b/Android.bp
@@ -51,6 +51,7 @@ bootstrap_go_package {
         "android/paths.go",
         "android/prebuilt.go",
         "android/register.go",
+        "android/testing.go",
         "android/util.go",
         "android/variable.go",
 
diff --git a/android/config.go b/android/config.go
index 3de1ef1fa..661e333d2 100644
--- a/android/config.go
+++ b/android/config.go
@@ -84,6 +84,8 @@ type config struct {
 
 	inMake bool
 
+	captureBuild bool // true for tests, saves build parameters for each module
+
 	OncePer
 }
 
@@ -171,7 +173,8 @@ func TestConfig(buildDir string) Config {
 			DeviceName: stringPtr("test_device"),
 		},
 
-		buildDir: buildDir,
+		buildDir:     buildDir,
+		captureBuild: true,
 	}
 	config.deviceConfig = &deviceConfig{
 		config: config,
diff --git a/android/defaults.go b/android/defaults.go
index 0776405fa..84f0a3dbd 100644
--- a/android/defaults.go
+++ b/android/defaults.go
@@ -113,6 +113,11 @@ func (defaultable *DefaultableModule) applyDefaults(ctx TopDownMutatorContext,
 	}
 }
 
+func registerDefaultsPreArchMutators(ctx RegisterMutatorsContext) {
+	ctx.BottomUp("defaults_deps", defaultsDepsMutator).Parallel()
+	ctx.TopDown("defaults", defaultsMutator).Parallel()
+}
+
 func defaultsDepsMutator(ctx BottomUpMutatorContext) {
 	if defaultable, ok := ctx.Module().(Defaultable); ok {
 		ctx.AddDependency(ctx.Module(), DefaultsDepTag, defaultable.defaults().Defaults...)
diff --git a/android/module.go b/android/module.go
index 7f541be29..fcb5aeefe 100644
--- a/android/module.go
+++ b/android/module.go
@@ -111,6 +111,8 @@ type Module interface {
 
 	AddProperties(props ...interface{})
 	GetProperties() []interface{}
+
+	BuildParamsForTests() []ModuleBuildParams
 }
 
 type nameProperties struct {
@@ -291,6 +293,9 @@ type ModuleBase struct {
 	hooks hooks
 
 	registerProps []interface{}
+
+	// For tests
+	buildParams []ModuleBuildParams
 }
 
 func (a *ModuleBase) AddProperties(props ...interface{}) {
@@ -301,6 +306,10 @@ func (a *ModuleBase) GetProperties() []interface{} {
 	return a.registerProps
 }
 
+func (a *ModuleBase) BuildParamsForTests() []ModuleBuildParams {
+	return a.buildParams
+}
+
 // Name returns the name of the module.  It may be overridden by individual module types, for
 // example prebuilts will prepend prebuilt_ to the name.
 func (a *ModuleBase) Name() string {
@@ -520,6 +529,8 @@ func (a *ModuleBase) GenerateBuildActions(ctx blueprint.ModuleContext) {
 			return
 		}
 	}
+
+	a.buildParams = androidCtx.buildParams
 }
 
 type androidBaseContextImpl struct {
@@ -538,6 +549,9 @@ type androidModuleContext struct {
 	checkbuildFiles Paths
 	missingDeps     []string
 	module          Module
+
+	// For tests
+	buildParams []ModuleBuildParams
 }
 
 func (a *androidModuleContext) ninjaError(desc string, outputs []string, err error) {
@@ -566,6 +580,10 @@ func (a *androidModuleContext) Build(pctx blueprint.PackageContext, params bluep
 }
 
 func (a *androidModuleContext) ModuleBuild(pctx blueprint.PackageContext, params ModuleBuildParams) {
+	if a.config.captureBuild {
+		a.buildParams = append(a.buildParams, params)
+	}
+
 	bparams := blueprint.BuildParams{
 		Rule:            params.Rule,
 		Deps:            params.Deps,
diff --git a/android/mutator.go b/android/mutator.go
index bb494878d..86a5b852a 100644
--- a/android/mutator.go
+++ b/android/mutator.go
@@ -15,8 +15,6 @@
 package android
 
 import (
-	"sync"
-
 	"github.com/google/blueprint"
 )
 
@@ -27,9 +25,6 @@ import (
 //   Deps
 //   PostDeps
 
-var registerMutatorsOnce sync.Once
-var registeredMutators []*mutator
-
 func registerMutatorsToContext(ctx *blueprint.Context, mutators []*mutator) {
 	for _, t := range mutators {
 		var handle blueprint.MutatorHandle
@@ -44,56 +39,24 @@ func registerMutatorsToContext(ctx *blueprint.Context, mutators []*mutator) {
 	}
 }
 
-func registerMutators(ctx *blueprint.Context) {
-
-	registerMutatorsOnce.Do(func() {
-		ctx := &registerMutatorsContext{}
-
-		register := func(funcs []RegisterMutatorFunc) {
-			for _, f := range funcs {
-				f(ctx)
-			}
-		}
-
-		ctx.TopDown("load_hooks", loadHookMutator).Parallel()
-		ctx.BottomUp("prebuilts", prebuiltMutator).Parallel()
-		ctx.BottomUp("defaults_deps", defaultsDepsMutator).Parallel()
-		ctx.TopDown("defaults", defaultsMutator).Parallel()
-
-		register(preArch)
-
-		ctx.BottomUp("arch", archMutator).Parallel()
-		ctx.TopDown("arch_hooks", archHookMutator).Parallel()
-
-		register(preDeps)
-
-		ctx.BottomUp("deps", depsMutator).Parallel()
-
-		ctx.TopDown("prebuilt_select", PrebuiltSelectModuleMutator).Parallel()
-		ctx.BottomUp("prebuilt_replace", PrebuiltReplaceMutator).Parallel()
-
-		register(postDeps)
-
-		registeredMutators = ctx.mutators
-	})
-
-	registerMutatorsToContext(ctx, registeredMutators)
-}
-
-func RegisterTestMutators(ctx *blueprint.Context) {
-	mutators := &registerMutatorsContext{}
+func registerMutators(ctx *blueprint.Context, preArch, preDeps, postDeps []RegisterMutatorFunc) {
+	mctx := &registerMutatorsContext{}
 
 	register := func(funcs []RegisterMutatorFunc) {
 		for _, f := range funcs {
-			f(mutators)
+			f(mctx)
 		}
 	}
 
-	register(testPreDeps)
-	mutators.BottomUp("deps", depsMutator).Parallel()
-	register(testPostDeps)
+	register(preArch)
 
-	registerMutatorsToContext(ctx, mutators.mutators)
+	register(preDeps)
+
+	mctx.BottomUp("deps", depsMutator).Parallel()
+
+	register(postDeps)
+
+	registerMutatorsToContext(ctx, mctx.mutators)
 }
 
 type registerMutatorsContext struct {
@@ -107,7 +70,24 @@ type RegisterMutatorsContext interface {
 
 type RegisterMutatorFunc func(RegisterMutatorsContext)
 
-var preArch, preDeps, postDeps, testPreDeps, testPostDeps []RegisterMutatorFunc
+var preArch = []RegisterMutatorFunc{
+	func(ctx RegisterMutatorsContext) {
+		ctx.TopDown("load_hooks", loadHookMutator).Parallel()
+	},
+	registerPrebuiltsPreArchMutators,
+	registerDefaultsPreArchMutators,
+}
+
+var preDeps = []RegisterMutatorFunc{
+	func(ctx RegisterMutatorsContext) {
+		ctx.BottomUp("arch", archMutator).Parallel()
+		ctx.TopDown("arch_hooks", archHookMutator).Parallel()
+	},
+}
+
+var postDeps = []RegisterMutatorFunc{
+	registerPrebuiltsPostDepsMutators,
+}
 
 func PreArchMutators(f RegisterMutatorFunc) {
 	preArch = append(preArch, f)
@@ -121,14 +101,6 @@ func PostDepsMutators(f RegisterMutatorFunc) {
 	postDeps = append(postDeps, f)
 }
 
-func TestPreDepsMutators(f RegisterMutatorFunc) {
-	testPreDeps = append(testPreDeps, f)
-}
-
-func TeststPostDepsMutators(f RegisterMutatorFunc) {
-	testPostDeps = append(testPostDeps, f)
-}
-
 type AndroidTopDownMutator func(TopDownMutatorContext)
 
 type TopDownMutatorContext interface {
diff --git a/android/prebuilt.go b/android/prebuilt.go
index 772df7af4..080df919d 100644
--- a/android/prebuilt.go
+++ b/android/prebuilt.go
@@ -61,6 +61,15 @@ type PrebuiltInterface interface {
 	Prebuilt() *Prebuilt
 }
 
+func registerPrebuiltsPreArchMutators(ctx RegisterMutatorsContext) {
+	ctx.BottomUp("prebuilts", prebuiltMutator).Parallel()
+}
+
+func registerPrebuiltsPostDepsMutators(ctx RegisterMutatorsContext) {
+	ctx.TopDown("prebuilt_select", PrebuiltSelectModuleMutator).Parallel()
+	ctx.BottomUp("prebuilt_replace", PrebuiltReplaceMutator).Parallel()
+}
+
 // prebuiltMutator ensures that there is always a module with an undecorated name, and marks
 // prebuilt modules that have both a prebuilt and a source module.
 func prebuiltMutator(ctx BottomUpMutatorContext) {
diff --git a/android/prebuilt_test.go b/android/prebuilt_test.go
index 5fa20326f..fe763eddc 100644
--- a/android/prebuilt_test.go
+++ b/android/prebuilt_test.go
@@ -122,9 +122,12 @@ func TestPrebuilts(t *testing.T) {
 
 	for _, test := range prebuiltsTests {
 		t.Run(test.name, func(t *testing.T) {
-			ctx := NewContext()
+			ctx := NewTestContext()
+			ctx.PreArchMutators(registerPrebuiltsPreArchMutators)
+			ctx.PostDepsMutators(registerPrebuiltsPostDepsMutators)
 			ctx.RegisterModuleType("prebuilt", ModuleFactoryAdaptor(newPrebuiltModule))
 			ctx.RegisterModuleType("source", ModuleFactoryAdaptor(newSourceModule))
+			ctx.Register()
 			ctx.MockFileSystem(map[string][]byte{
 				"Blueprints": []byte(`
 					source {
@@ -139,13 +142,10 @@ func TestPrebuilts(t *testing.T) {
 			_, errs = ctx.PrepareBuildActions(config)
 			fail(t, errs)
 
-			foo := findModule(ctx, "foo")
-			if foo == nil {
-				t.Fatalf("failed to find module foo")
-			}
+			foo := ctx.ModuleForTests("foo", "")
 
 			var dependsOnSourceModule, dependsOnPrebuiltModule bool
-			ctx.VisitDirectDeps(foo, func(m blueprint.Module) {
+			ctx.VisitDirectDeps(foo.Module(), func(m blueprint.Module) {
 				if _, ok := m.(*sourceModule); ok {
 					dependsOnSourceModule = true
 				}
@@ -228,16 +228,6 @@ func (s *sourceModule) DepsMutator(ctx BottomUpMutatorContext) {
 func (s *sourceModule) GenerateAndroidBuildActions(ctx ModuleContext) {
 }
 
-func findModule(ctx *blueprint.Context, name string) blueprint.Module {
-	var ret blueprint.Module
-	ctx.VisitAllModules(func(m blueprint.Module) {
-		if ctx.ModuleName(m) == name {
-			ret = m
-		}
-	})
-	return ret
-}
-
 func fail(t *testing.T, errs []error) {
 	if len(errs) > 0 {
 		for _, err := range errs {
diff --git a/android/register.go b/android/register.go
index 76a1cc9f7..226e790cd 100644
--- a/android/register.go
+++ b/android/register.go
@@ -60,9 +60,15 @@ func RegisterSingletonType(name string, factory blueprint.SingletonFactory) {
 	singletons = append(singletons, singleton{name, factory})
 }
 
-func NewContext() *blueprint.Context {
-	ctx := blueprint.NewContext()
+type Context struct {
+	*blueprint.Context
+}
 
+func NewContext() *Context {
+	return &Context{blueprint.NewContext()}
+}
+
+func (ctx *Context) Register() {
 	for _, t := range moduleTypes {
 		ctx.RegisterModuleType(t.name, t.factory)
 	}
@@ -71,9 +77,7 @@ func NewContext() *blueprint.Context {
 		ctx.RegisterSingletonType(t.name, t.factory)
 	}
 
-	registerMutators(ctx)
+	registerMutators(ctx.Context, preArch, preDeps, postDeps)
 
 	ctx.RegisterSingletonType("env", EnvSingleton)
-
-	return ctx
 }
diff --git a/android/testing.go b/android/testing.go
new file mode 100644
index 000000000..414477510
--- /dev/null
+++ b/android/testing.go
@@ -0,0 +1,98 @@
+// Copyright 2017 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package android
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/google/blueprint"
+)
+
+func NewTestContext() *TestContext {
+	return &TestContext{
+		Context: blueprint.NewContext(),
+	}
+}
+
+type TestContext struct {
+	*blueprint.Context
+	preArch, preDeps, postDeps []RegisterMutatorFunc
+}
+
+func (ctx *TestContext) PreArchMutators(f RegisterMutatorFunc) {
+	ctx.preArch = append(ctx.preArch, f)
+}
+
+func (ctx *TestContext) PreDepsMutators(f RegisterMutatorFunc) {
+	ctx.preDeps = append(ctx.preDeps, f)
+}
+
+func (ctx *TestContext) PostDepsMutators(f RegisterMutatorFunc) {
+	ctx.postDeps = append(ctx.postDeps, f)
+}
+
+func (ctx *TestContext) Register() {
+	registerMutators(ctx.Context, ctx.preArch, ctx.preDeps, ctx.postDeps)
+
+	ctx.RegisterSingletonType("env", EnvSingleton)
+}
+
+func (ctx *TestContext) ModuleForTests(name, variant string) TestingModule {
+	var module Module
+	ctx.VisitAllModules(func(m blueprint.Module) {
+		if ctx.ModuleName(m) == name && ctx.ModuleSubDir(m) == variant {
+			module = m.(Module)
+		}
+	})
+
+	if module == nil {
+		panic(fmt.Errorf("failed to find module %q variant %q", name, variant))
+	}
+
+	return TestingModule{module}
+}
+
+type TestingModule struct {
+	module Module
+}
+
+func (m TestingModule) Module() Module {
+	return m.module
+}
+
+func (m TestingModule) Rule(rule string) ModuleBuildParams {
+	for _, p := range m.module.BuildParamsForTests() {
+		if strings.Contains(p.Rule.String(), rule) {
+			return p
+		}
+	}
+	panic(fmt.Errorf("couldn't find rule %q", rule))
+}
+
+func (m TestingModule) Output(file string) ModuleBuildParams {
+	for _, p := range m.module.BuildParamsForTests() {
+		outputs := append(WritablePaths(nil), p.Outputs...)
+		if p.Output != nil {
+			outputs = append(outputs, p.Output)
+		}
+		for _, f := range outputs {
+			if f.Base() == file {
+				return p
+			}
+		}
+	}
+	panic(fmt.Errorf("couldn't find output %q", file))
+}
diff --git a/cc/test_data_test.go b/cc/test_data_test.go
index a798919f8..962bde5bc 100644
--- a/cc/test_data_test.go
+++ b/cc/test_data_test.go
@@ -23,8 +23,6 @@ import (
 
 	"android/soong/android"
 	"android/soong/genrule"
-
-	"github.com/google/blueprint"
 )
 
 type dataFile struct {
@@ -123,8 +121,7 @@ func TestDataTests(t *testing.T) {
 
 	for _, test := range testDataTests {
 		t.Run(test.name, func(t *testing.T) {
-			ctx := blueprint.NewContext()
-			android.RegisterTestMutators(ctx)
+			ctx := android.NewTestContext()
 			ctx.MockFileSystem(map[string][]byte{
 				"Blueprints":     []byte(`subdirs = ["dir"]`),
 				"dir/Blueprints": []byte(test.modules),
@@ -135,18 +132,16 @@ func TestDataTests(t *testing.T) {
 				android.ModuleFactoryAdaptor(genrule.FileGroupFactory))
 			ctx.RegisterModuleType("test",
 				android.ModuleFactoryAdaptor(newTest))
+			ctx.Register()
 
 			_, errs := ctx.ParseBlueprintsFiles("Blueprints")
 			fail(t, errs)
 			_, errs = ctx.PrepareBuildActions(config)
 			fail(t, errs)
 
-			foo := findModule(ctx, "foo")
-			if foo == nil {
-				t.Fatalf("failed to find module foo")
-			}
+			foo := ctx.ModuleForTests("foo", "")
 
-			got := foo.(*testDataTest).data
+			got := foo.Module().(*testDataTest).data
 			if len(got) != len(test.data) {
 				t.Errorf("expected %d data files, got %d",
 					len(test.data), len(got))
@@ -192,16 +187,6 @@ func (test *testDataTest) GenerateAndroidBuildActions(ctx android.ModuleContext)
 	test.data = ctx.ExpandSources(test.Properties.Data, nil)
 }
 
-func findModule(ctx *blueprint.Context, name string) blueprint.Module {
-	var ret blueprint.Module
-	ctx.VisitAllModules(func(m blueprint.Module) {
-		if ctx.ModuleName(m) == name {
-			ret = m
-		}
-	})
-	return ret
-}
-
 func fail(t *testing.T, errs []error) {
 	if len(errs) > 0 {
 		for _, err := range errs {
diff --git a/cmd/soong_build/main.go b/cmd/soong_build/main.go
index ed8f2fd58..e15a6bdd1 100644
--- a/cmd/soong_build/main.go
+++ b/cmd/soong_build/main.go
@@ -32,6 +32,7 @@ func main() {
 	srcDir := filepath.Dir(flag.Arg(0))
 
 	ctx := android.NewContext()
+	ctx.Register()
 
 	configuration, err := android.NewConfig(srcDir, bootstrap.BuildDir)
 	if err != nil {
@@ -44,5 +45,5 @@ func main() {
 
 	ctx.SetAllowMissingDependencies(configuration.AllowMissingDependencies())
 
-	bootstrap.Main(ctx, configuration, configuration.ConfigFileName, configuration.ProductVariablesFileName)
+	bootstrap.Main(ctx.Context, configuration, configuration.ConfigFileName, configuration.ProductVariablesFileName)
 }
diff --git a/python/python_test.go b/python/python_test.go
index 57aaa344d..4c30d95be 100644
--- a/python/python_test.go
+++ b/python/python_test.go
@@ -26,8 +26,6 @@ import (
 	"testing"
 
 	"android/soong/android"
-
-	"github.com/google/blueprint"
 )
 
 type pyBinary struct {
@@ -306,17 +304,17 @@ var (
 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 @@ func expectErrors(t *testing.T, actErrs []error, expErrs []string) (testErrs []e
 	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 @@ func tearDownBuildEnv(buildDir string) {
 	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 {