diff --git a/android/bazel_handler.go b/android/bazel_handler.go
index 28c0e53..7911632 100644
--- a/android/bazel_handler.go
+++ b/android/bazel_handler.go
@@ -117,21 +117,25 @@
 
 // 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 @@
 	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{}