binary search tool: Provide temp files that hold GOOD/BAD sets

Provide temp files to scripts ($BISECT_GOOD_SET, $BISECT_BAD_SET) that
provide a full listing of the items in the good set and items in the bad
set. This can be used for debugging purposes or so scripts can analyze
the state of the current binary search iteration.

TEST=Add unit test that utilizes these files

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

4 files changed, 142 insertions, 23 deletions

diff --git a/binary_search_tool/binary_search_state.py b/binary_search_tool/binary_search_state.py
index 1f797281..55d39d6d 100755
--- a/binary_search_tool/binary_search_state.py
+++ b/binary_search_tool/binary_search_state.py
@@ -4,6 +4,7 @@
 from __future__ import print_function
 
 import argparse
+import contextlib
 import math
 import os
 import pickle
@@ -20,6 +21,9 @@ from cros_utils import logger
 
 import binary_search_perforce
 
+GOOD_SET_VAR = 'BISECT_GOOD_SET'
+BAD_SET_VAR = 'BISECT_BAD_SET'
+
 STATE_FILE = '%s.state' % sys.argv[0]
 HIDDEN_STATE_FILE = os.path.join(
     os.path.dirname(STATE_FILE), '.%s' % os.path.basename(STATE_FILE))
@@ -29,6 +33,30 @@ class Error(Exception):
   pass
 
 
+@contextlib.contextmanager
+def SetFile(env_var, items):
+  """Generate set files that can be used by switch/test scripts.
+
+  Generate temporary set file (good/bad) holding contents of good/bad items for
+  the current binary search iteration. Store the name of each file as an
+  environment variable so all child processes can access it.
+
+  This function is a contextmanager, meaning it's meant to be used with the
+  "with" statement in Python. This is so cleanup and setup happens automatically
+  and cleanly. Execution of the outer "with" statement happens at the "yield"
+  statement.
+
+  Args:
+    env_var: What environment variable to store the file name in.
+    items: What items are in this set.
+  """
+  with tempfile.NamedTemporaryFile() as f:
+    os.environ[env_var] = f.name
+    f.write('\n'.join(items))
+    f.flush()
+    yield
+
+
 class BinarySearchState(object):
   """The binary search state class."""
 
@@ -155,19 +183,21 @@ class BinarySearchState(object):
     self.l.LogOutput('Beginning %d tests to verify good/bad sets\n' %
                      self.verify_level)
     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.'
+      with SetFile(GOOD_SET_VAR, self.all_items), SetFile(BAD_SET_VAR, []):
+        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.'
+
+      with SetFile(GOOD_SET_VAR, []), SetFile(BAD_SET_VAR, self.all_items):
+        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.'
 
   def DoSearch(self):
     """Perform full search for bad items.
@@ -226,16 +256,17 @@ class BinarySearchState(object):
       self.search_cycles += 1
       [bad_items, good_items] = self.GetNextItems()
 
-      # TODO: bad_items should come first.
-      self.SwitchToGood(good_items)
-      self.SwitchToBad(bad_items)
-      status = self.InstallScript()
-      if status == 0:
-        status = self.TestScript()
-      else:
-        # Install script failed, treat as skipped item
-        status = 2
-      terminated = self.binary_search.SetStatus(status)
+      with SetFile(GOOD_SET_VAR, good_items), SetFile(BAD_SET_VAR, bad_items):
+        # TODO: bad_items should come first.
+        self.SwitchToGood(good_items)
+        self.SwitchToBad(bad_items)
+        status = self.InstallScript()
+        if status == 0:
+          status = self.TestScript()
+        else:
+          # Install script failed, treat as skipped item
+          status = 2
+        terminated = self.binary_search.SetStatus(status)
 
       if terminated:
         self.l.LogOutput('Terminated!')
diff --git a/binary_search_tool/test/binary_search_tool_tester.py b/binary_search_tool/test/binary_search_tool_tester.py
index 23cc0096..dc468e86 100755
--- a/binary_search_tool/test/binary_search_tool_tester.py
+++ b/binary_search_tool/test/binary_search_tool_tester.py
@@ -285,6 +285,17 @@ class BisectingUtilsTest(unittest.TestCase):
     found_obj = int(bss.found_items.pop())
     self.assertEquals(bad_objs[found_obj], 1)
 
+  def test_set_file(self):
+    binary_search_state.Run(
+        get_initial_items='./gen_init_list.py',
+        switch_to_good='./switch_to_good_set_file.py',
+        switch_to_bad='./switch_to_bad_set_file.py',
+        test_script='./is_good.py',
+        prune=True,
+        file_args=True,
+        verify_level=1)
+    self.check_output()
+
   def check_output(self):
     _, out, _ = command_executer.GetCommandExecuter().RunCommandWOutput(
         ('grep "Bad items are: " logs/binary_search_tool_tester.py.out | '
@@ -323,6 +334,7 @@ def Main(argv):
   suite.addTest(BisectingUtilsTest('test_verify_fail'))
   suite.addTest(BisectingUtilsTest('test_early_terminate'))
   suite.addTest(BisectingUtilsTest('test_no_prune'))
+  suite.addTest(BisectingUtilsTest('test_set_file'))
   suite.addTest(BisectTest('test_full_bisector'))
   runner = unittest.TextTestRunner()
   runner.run(suite)
diff --git a/binary_search_tool/test/switch_to_bad_set_file.py b/binary_search_tool/test/switch_to_bad_set_file.py
new file mode 100755
index 00000000..f535fdfd
--- /dev/null
+++ b/binary_search_tool/test/switch_to_bad_set_file.py
@@ -0,0 +1,37 @@
+#!/usr/bin/python2
+"""Switch part of the objects file in working set to (possible) bad ones.
+
+This script is meant to be specifically used with the set_file test. This uses
+the set files generated by binary_search_state to do the switching.
+"""
+
+from __future__ import print_function
+
+import os
+import sys
+
+import common
+
+
+def Main(_):
+  """Switch part of the objects file in working set to (possible) bad ones."""
+  working_set = common.ReadWorkingSet()
+  objects_file = common.ReadObjectsFile()
+
+  if not os.path.exists(os.environ['BISECT_BAD_SET']):
+    print('Bad set file does not exist!')
+    return 1
+
+  object_index = common.ReadObjectIndex(os.environ['BISECT_BAD_SET'])
+
+  for oi in object_index:
+    working_set[int(oi)] = objects_file[oi]
+
+  common.WriteWorkingSet(working_set)
+
+  return 0
+
+
+if __name__ == '__main__':
+  retval = Main(sys.argv)
+  sys.exit(retval)
diff --git a/binary_search_tool/test/switch_to_good_set_file.py b/binary_search_tool/test/switch_to_good_set_file.py
new file mode 100755
index 00000000..83777af0
--- /dev/null
+++ b/binary_search_tool/test/switch_to_good_set_file.py
@@ -0,0 +1,39 @@
+#!/usr/bin/python2
+"""Change portions of the object files to good.
+
+The "portion" is defined by the file (which is passed as the only argument to
+this script) content. Every line in the file is an object index, which will be
+set to good (mark as 0).
+
+This script is meant to be specifically used with the set_file test. This uses
+the set files generated by binary_search_state to do the switching.
+"""
+
+from __future__ import print_function
+
+import os
+import sys
+
+import common
+
+
+def Main(_):
+  working_set = common.ReadWorkingSet()
+
+  if not os.path.exists(os.environ['BISECT_GOOD_SET']):
+    print('Good set file does not exist!')
+    return 1
+
+  object_index = common.ReadObjectIndex(os.environ['BISECT_GOOD_SET'])
+
+  for oi in object_index:
+    working_set[int(oi)] = 0
+
+  common.WriteWorkingSet(working_set)
+
+  return 0
+
+
+if __name__ == '__main__':
+  retval = Main(sys.argv)
+  sys.exit(retval)