Add install_script argument to binary search tool.

Add optional install_script argument that will run before test_script
use this install_script to perform building/flashing/other misc. setup.
If your build or flash succeeds then return 0, otherwise return 1 and
this build will be "skipped" (equivalent to the test script returning
a 2 exit code).

Test=Update unit tests to test for successful and failed install

Change-Id: Ic492f08230796f50b7bb93aebe078960419a3c99
Reviewed-on: https://chrome-internal-review.googlesource.com/262988
Commit-Ready: Cassidy Burden <cburden@google.com>
Tested-by: Cassidy Burden <cburden@google.com>
Reviewed-by: Caroline Tice <cmtice@google.com>

6 files changed, 91 insertions, 6 deletions

diff --git a/binary_search_tool/binary_search_state.py b/binary_search_tool/binary_search_state.py
index b8898c23..56eaee10 100755
--- a/binary_search_tool/binary_search_state.py
+++ b/binary_search_tool/binary_search_state.py
@@ -1,5 +1,4 @@
 #!/usr/bin/python2
-
 """The binary search wrapper."""
 
 from __future__ import print_function
@@ -30,12 +29,14 @@ STATE_FILE = '%s.state' % sys.argv[0]
 
 class BinarySearchState(object):
   """The binary search state class."""
+
   def __init__(self, get_initial_items, switch_to_good, switch_to_bad,
-               test_script, incremental, prune, iterations, prune_iterations,
-               verify_level, file_args):
+               install_script, test_script, incremental, prune, iterations,
+               prune_iterations, verify_level, file_args):
     self.get_initial_items = get_initial_items
     self.switch_to_good = switch_to_good
     self.switch_to_bad = switch_to_bad
+    self.install_script = install_script
     self.test_script = test_script
     self.incremental = incremental
     self.prune = prune
@@ -107,15 +108,26 @@ class BinarySearchState(object):
     command = self.test_script
     return self.ce.RunCommand(command)
 
+  def InstallScript(self):
+    if not self.install_script:
+      return 0
+
+    command = self.install_script
+    return self.ce.RunCommand(command)
+
   def DoVerify(self):
     for _ in range(int(self.verify_level)):
       self.l.LogOutput('Resetting all items to good to verify.')
       self.SwitchToGood(self.all_items)
+      status = self.InstallScript()
+      assert status == 0, 'When reset_to_good, install should succeed.'
       status = self.TestScript()
       assert status == 0, 'When reset_to_good, status should be 0.'
 
       self.l.LogOutput('Resetting all items to bad to verify.')
       self.SwitchToBad(self.all_items)
+      status = self.InstallScript()
+      assert status == 0, 'When reset_to_bad, install should succeed.'
       status = self.TestScript()
       assert status == 1, 'When reset_to_bad, status should be 1.'
 
@@ -173,7 +185,12 @@ class BinarySearchState(object):
       # TODO: bad_items should come first.
       self.SwitchToGood(good_items)
       self.SwitchToBad(bad_items)
-      status = self.TestScript()
+      status = self.InstallScript()
+      if status == 0:
+        status = self.TestScript()
+      else:
+        # Install script failed, treat as skipped item
+        status = 2
       terminated = self.bs.SetStatus(status)
 
       if terminated:
@@ -255,6 +272,12 @@ def Main(argv):
                       '--switch_to_bad',
                       dest='switch_to_bad',
                       help='Script to run to switch to bad.')
+  parser.add_argument('-I',
+                      '--install_script',
+                      dest='install_script',
+                      default=None,
+                      help=('Optional script to perform building, flashing, '
+                            'and other setup before the test script runs.'))
   parser.add_argument('-t',
                       '--test_script',
                       dest='test_script',
@@ -302,6 +325,9 @@ def Main(argv):
   iterations = int(options.iterations)
   switch_to_good = _CanonicalizeScript(options.switch_to_good)
   switch_to_bad = _CanonicalizeScript(options.switch_to_bad)
+  install_script = options.install_script
+  if install_script:
+    install_script = _CanonicalizeScript(options.install_script)
   test_script = _CanonicalizeScript(options.test_script)
   get_initial_items = _CanonicalizeScript(options.get_initial_items)
   prune = options.prune
@@ -316,8 +342,9 @@ def Main(argv):
 
   try:
     bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad,
-                            test_script, incremental, prune, iterations,
-                            prune_iterations, verify_level, file_args)
+                            install_script, test_script, incremental, prune,
+                            iterations, prune_iterations, verify_level,
+                            file_args)
     bss.DoVerify()
     bss.DoSearch()
 
diff --git a/binary_search_tool/test/binary_search_tool_tester.py b/binary_search_tool/test/binary_search_tool_tester.py
index 45ce0642..475680c6 100755
--- a/binary_search_tool/test/binary_search_tool_tester.py
+++ b/binary_search_tool/test/binary_search_tool_tester.py
@@ -42,7 +42,26 @@ class BisectingUtilsTest(unittest.TestCase):
             './switch_to_good.py', '--switch_to_bad', './switch_to_bad.py',
             '--test_script', './is_good.py', '--prune', '--file_args']
     binary_search_state.Main(args)
+    self.check_output()
 
+  def test_install_script(self):
+    args = ['--get_initial_items', './gen_init_list.py', '--switch_to_good',
+            './switch_to_good.py', '--switch_to_bad', './switch_to_bad.py',
+            '--test_script', './is_good.py', '--prune', '--file_args',
+            '--install_script', './install.py']
+    common.installed = False
+    binary_search_state.Main(args)
+    self.check_output()
+
+  def test_bad_install_script(self):
+    args = ['--get_initial_items', './gen_init_list.py', '--switch_to_good',
+            './switch_to_good.py', '--switch_to_bad', './switch_to_bad.py',
+            '--test_script', './is_good.py', '--prune', '--file_args',
+            '--install_script', './install_bad.py']
+    with self.assertRaises(AssertionError):
+      binary_search_state.Main(args)
+
+  def check_output(self):
     _, out, _ = command_executer.GetCommandExecuter().RunCommandWOutput(
         'tail -n 10 logs/binary_search_state.py.out')
     ls = out.splitlines()
@@ -63,6 +82,8 @@ def Main(argv):
   suite = unittest.TestSuite()
   for _ in range(0, num_tests):
     suite.addTest(BisectingUtilsTest())
+  suite.addTest(BisectingUtilsTest('test_install_script'))
+  suite.addTest(BisectingUtilsTest('test_bad_install_script'))
   runner = unittest.TextTestRunner()
   runner.run(suite)
 
diff --git a/binary_search_tool/test/common.py b/binary_search_tool/test/common.py
index baac9434..382274b9 100755
--- a/binary_search_tool/test/common.py
+++ b/binary_search_tool/test/common.py
@@ -6,6 +6,8 @@ DEFAULT_BAD_OBJECT_NUMBER = 23
 OBJECTS_FILE = 'objects.txt'
 WORKING_SET_FILE = 'working_set.txt'
 
+installed = True
+
 
 def ReadWorkingSet():
   working_set = []
diff --git a/binary_search_tool/test/install.py b/binary_search_tool/test/install.py
new file mode 100755
index 00000000..21cf97a1
--- /dev/null
+++ b/binary_search_tool/test/install.py
@@ -0,0 +1,18 @@
+#!/usr/bin/python2
+"""Emulate installation of files, is_good.py should fail without this."""
+
+from __future__ import print_function
+
+import sys
+
+import common
+
+
+def Main():
+  common.installed = True
+  return 0
+
+
+if __name__ == '__main__':
+  retval = Main()
+  sys.exit(retval)
diff --git a/binary_search_tool/test/install_bad.py b/binary_search_tool/test/install_bad.py
new file mode 100755
index 00000000..8b49a3f8
--- /dev/null
+++ b/binary_search_tool/test/install_bad.py
@@ -0,0 +1,15 @@
+#!/usr/bin/python2
+"""Emulate installation that fails (i.e. failed flash to device)"""
+
+from __future__ import print_function
+
+import sys
+
+
+def Main():
+  return 1  ## False, flashing failure
+
+
+if __name__ == '__main__':
+  retval = Main()
+  sys.exit(retval)
diff --git a/binary_search_tool/test/is_good.py b/binary_search_tool/test/is_good.py
index 586d6911..e2ab4642 100755
--- a/binary_search_tool/test/is_good.py
+++ b/binary_search_tool/test/is_good.py
@@ -9,6 +9,8 @@ import common
 
 
 def Main():
+  if not common.installed:
+    return 1
   working_set = common.ReadWorkingSet()
   for w in working_set:
     if w == 1: