Merge "Fix broken test_lunch.py"
diff --git a/orchestrator/core/lunch.py b/orchestrator/core/lunch.py
index a648478..70a2d1d 100755
--- a/orchestrator/core/lunch.py
+++ b/orchestrator/core/lunch.py
@@ -75,11 +75,17 @@
     for f in walk_paths(path, lambda x: x == filename):
         return f
 
+# TODO: When orchestrator is in its own git project remove the "build" and "make" here
+class LunchContext(object):
+    """Mockable container for lunch"""
+    def __init__(self, workspace_root, orchestrator_path_prefix_components=["build", "build", "make"]):
+      self.workspace_root = workspace_root
+      self.orchestrator_path_prefix_components = orchestrator_path_prefix_components
 
-def find_config_dirs(workspace_root):
+def find_config_dirs(context):
     """Find the configuration files in the well known locations inside workspace_root
 
-        <workspace_root>/build/build/orchestrator/multitree_combos
+        <workspace_root>/<orchestrator>/<path>/<prefix>/orchestrator/multitree_combos
            (AOSP devices, such as cuttlefish)
 
         <workspace_root>/vendor/**/multitree_combos
@@ -93,21 +99,20 @@
     """
     # TODO: This is not looking in inner trees correctly.
 
-    # TODO: When orchestrator is in its own git project remove the "make/" here
-    yield os.path.join(workspace_root, "build/build/make/orchestrator/multitree_combos")
+    yield os.path.join(context.workspace_root, *context.orchestrator_path_prefix_components, "orchestrator/multitree_combos")
 
     dirs = ["vendor", "device"]
     for d in dirs:
-        yield from find_dirs(os.path.join(workspace_root, d), "multitree_combos")
+        yield from find_dirs(os.path.join(context.workspace_root, d), "multitree_combos")
 
 
-def find_named_config(workspace_root, shortname):
-    """Find the config with the given shortname inside workspace_root.
+def find_named_config(context, shortname):
+    """Find the config with the given shortname inside context.workspace_root.
 
     Config directories are searched in the order described in find_config_dirs,
     and inside those directories, alphabetically."""
     filename = shortname + ".mcombo"
-    for config_dir in find_config_dirs(workspace_root):
+    for config_dir in find_config_dirs(context):
         found = find_file(config_dir, filename)
         if found:
             return found
@@ -122,7 +127,7 @@
     return split
 
 
-def choose_config_from_args(workspace_root, args):
+def choose_config_from_args(context, args):
     """Return the config file we should use for the given argument,
     or null if there's no file that matches that."""
     if len(args) == 1:
@@ -130,7 +135,7 @@
         # file we don't match that.
         pv = parse_product_variant(args[0])
         if pv:
-            config = find_named_config(workspace_root, pv[0])
+            config = find_named_config(context, pv[0])
             if config:
                 return (config, pv[1])
             return None, None
@@ -295,9 +300,9 @@
     return EXIT_STATUS_OK
 
 
-def find_all_combo_files(workspace_root):
+def find_all_combo_files(context):
     """Find all .mcombo files in the prescribed locations in the tree."""
-    for dir in find_config_dirs(workspace_root):
+    for dir in find_config_dirs(context):
         for file in walk_paths(dir, lambda x: x.endswith(".mcombo")):
             yield file
 
@@ -313,10 +318,10 @@
     return config.get("lunchable", False)
 
 
-def find_all_lunchable(workspace_root):
-    """Find all mcombo files in the tree (rooted at workspace_root) that when
+def find_all_lunchable(context):
+    """Find all mcombo files in the tree (rooted at context.workspace_root) that when
     parsed (and inheritance is flattened) have lunchable: true."""
-    for f in [x for x in find_all_combo_files(workspace_root) if is_file_lunchable(x)]:
+    for f in [x for x in find_all_combo_files(context) if is_file_lunchable(x)]:
         yield f
 
 
@@ -353,7 +358,8 @@
 
 def do_list():
     """Handle the --list command."""
-    for f in sorted(find_all_lunchable(".")):
+    lunch_context = LunchContext(".")
+    for f in sorted(find_all_lunchable(lunch_context)):
         print(f)
 
 
diff --git a/orchestrator/core/test_lunch.py b/orchestrator/core/test_lunch.py
index 2d85d05..5b890fb 100755
--- a/orchestrator/core/test_lunch.py
+++ b/orchestrator/core/test_lunch.py
@@ -20,6 +20,11 @@
 sys.dont_write_bytecode = True
 import lunch
 
+# Create a test LunchContext object
+# Test workspace is in test/configs
+# Orchestrator prefix inside it is build/make
+test_lunch_context = lunch.LunchContext("test/configs", ["build", "make"])
+
 class TestStringMethods(unittest.TestCase):
 
     def test_find_dirs(self):
@@ -35,61 +40,61 @@
                    "test/configs/device/aa/bb/multitree_combos/v.mcombo")
 
     def test_find_config_dirs(self):
-        self.assertEqual([x for x in lunch.find_config_dirs("test/configs")], [
+        self.assertEqual([x for x in lunch.find_config_dirs(test_lunch_context)], [
                     "test/configs/build/make/orchestrator/multitree_combos",
                     "test/configs/vendor/aa/bb/multitree_combos",
                     "test/configs/device/aa/bb/multitree_combos"])
 
     def test_find_named_config(self):
         # Inside build/orchestrator, overriding device and vendor
-        self.assertEqual(lunch.find_named_config("test/configs", "b"),
+        self.assertEqual(lunch.find_named_config(test_lunch_context, "b"),
                     "test/configs/build/make/orchestrator/multitree_combos/b.mcombo")
 
         # Nested dir inside a combo dir
-        self.assertEqual(lunch.find_named_config("test/configs", "nested"),
+        self.assertEqual(lunch.find_named_config(test_lunch_context, "nested"),
                     "test/configs/build/make/orchestrator/multitree_combos/nested/nested.mcombo")
 
         # Inside vendor, overriding device
-        self.assertEqual(lunch.find_named_config("test/configs", "v"),
+        self.assertEqual(lunch.find_named_config(test_lunch_context, "v"),
                     "test/configs/vendor/aa/bb/multitree_combos/v.mcombo")
 
         # Inside device
-        self.assertEqual(lunch.find_named_config("test/configs", "d"),
+        self.assertEqual(lunch.find_named_config(test_lunch_context, "d"),
                     "test/configs/device/aa/bb/multitree_combos/d.mcombo")
 
         # Make sure we don't look too deep (for performance)
-        self.assertIsNone(lunch.find_named_config("test/configs", "too_deep"))
+        self.assertIsNone(lunch.find_named_config(test_lunch_context, "too_deep"))
 
 
     def test_choose_config_file(self):
         # Empty string argument
-        self.assertEqual(lunch.choose_config_from_args("test/configs", [""]),
+        self.assertEqual(lunch.choose_config_from_args(test_lunch_context, [""]),
                     (None, None))
 
         # A PRODUCT-VARIANT name
-        self.assertEqual(lunch.choose_config_from_args("test/configs", ["v-eng"]),
+        self.assertEqual(lunch.choose_config_from_args(test_lunch_context, ["v-eng"]),
                     ("test/configs/vendor/aa/bb/multitree_combos/v.mcombo", "eng"))
 
         # A PRODUCT-VARIANT name that conflicts with a file
-        self.assertEqual(lunch.choose_config_from_args("test/configs", ["b-eng"]),
+        self.assertEqual(lunch.choose_config_from_args(test_lunch_context, ["b-eng"]),
                     ("test/configs/build/make/orchestrator/multitree_combos/b.mcombo", "eng"))
 
         # A PRODUCT-VARIANT that doesn't exist
-        self.assertEqual(lunch.choose_config_from_args("test/configs", ["z-user"]),
+        self.assertEqual(lunch.choose_config_from_args(test_lunch_context, ["z-user"]),
                     (None, None))
 
         # An explicit file
-        self.assertEqual(lunch.choose_config_from_args("test/configs",
+        self.assertEqual(lunch.choose_config_from_args(test_lunch_context,
                         ["test/configs/build/make/orchestrator/multitree_combos/b.mcombo", "eng"]),
                     ("test/configs/build/make/orchestrator/multitree_combos/b.mcombo", "eng"))
 
         # An explicit file that doesn't exist
-        self.assertEqual(lunch.choose_config_from_args("test/configs",
+        self.assertEqual(lunch.choose_config_from_args(test_lunch_context,
                         ["test/configs/doesnt_exist.mcombo", "eng"]),
                     (None, None))
 
         # An explicit file without a variant should fail
-        self.assertEqual(lunch.choose_config_from_args("test/configs",
+        self.assertEqual(lunch.choose_config_from_args(test_lunch_context,
                         ["test/configs/build/make/orchestrator/multitree_combos/b.mcombo"]),
                     ("test/configs/build/make/orchestrator/multitree_combos/b.mcombo", None))
 
@@ -119,7 +124,7 @@
                         })
 
     def test_list(self):
-        self.assertEqual(sorted(lunch.find_all_lunchable("test/configs")),
+        self.assertEqual(sorted(lunch.find_all_lunchable(test_lunch_context)),
                 ["test/configs/build/make/orchestrator/multitree_combos/b.mcombo"])
 
 if __name__ == "__main__":