diff --git a/android/bazel_handler.go b/android/bazel_handler.go
index 28c0e5307..7911632b6 100644
--- a/android/bazel_handler.go
+++ b/android/bazel_handler.go
@@ -117,21 +117,25 @@ var _ BazelContext = noopBazelContext{}
 
 // A bazel context to use for tests.
 type MockBazelContext struct {
-	AllFiles map[string][]string
+	OutputBaseDir string
+
+	LabelToOutputFiles                 map[string][]string
+	LabelToOutputFilesAndCcObjectFiles map[string]cquery.GetOutputFilesAndCcObjectFiles_Result
+	LabelToCcStaticLibraryFiles        map[string][]string
 }
 
 func (m MockBazelContext) GetOutputFiles(label string, archType ArchType) ([]string, bool) {
-	result, ok := m.AllFiles[label]
+	result, ok := m.LabelToOutputFiles[label]
 	return result, ok
 }
 
 func (m MockBazelContext) GetOutputFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool) {
-	result, ok := m.AllFiles[label]
-	return result, result, ok
+	result, ok := m.LabelToOutputFilesAndCcObjectFiles[label]
+	return result.OutputFiles, result.CcObjectFiles, ok
 }
 
 func (m MockBazelContext) GetPrebuiltCcStaticLibraryFiles(label string, archType ArchType) ([]string, bool) {
-	result, ok := m.AllFiles[label]
+	result, ok := m.LabelToCcStaticLibraryFiles[label]
 	return result, ok
 }
 
@@ -143,9 +147,7 @@ func (m MockBazelContext) BazelEnabled() bool {
 	return true
 }
 
-func (m MockBazelContext) OutputBase() string {
-	return "outputbase"
-}
+func (m MockBazelContext) OutputBase() string { return m.OutputBaseDir }
 
 func (m MockBazelContext) BuildStatementsToRegister() []bazel.BuildStatement {
 	return []bazel.BuildStatement{}
diff --git a/genrule/genrule_test.go b/genrule/genrule_test.go
index 3f1e9f3b9..3ce4f85f9 100644
--- a/genrule/genrule_test.go
+++ b/genrule/genrule_test.go
@@ -696,7 +696,8 @@ func TestGenruleWithBazel(t *testing.T) {
 	result := android.GroupFixturePreparers(
 		prepareForGenRuleTest, android.FixtureModifyConfig(func(config android.Config) {
 			config.BazelContext = android.MockBazelContext{
-				AllFiles: map[string][]string{
+				OutputBaseDir: "outputbase",
+				LabelToOutputFiles: map[string][]string{
 					"//foo/bar:bar": []string{"bazelone.txt", "bazeltwo.txt"}}}
 		})).RunTestWithBp(t, testGenruleBp()+bp)