Fix binary search tool's temp file creation

Fix the binary search tool so that temp files are created properly. This
means that temp files made for the switch scripts should be automatically
deleted once the switch script completes.

TEST=New unit test for temp file deletion

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

2 files changed, 51 insertions, 0 deletions

diff --git a/binary_search_tool/test/binary_search_tool_tester.py b/binary_search_tool/test/binary_search_tool_tester.py
index 451543d1..0f7906b0 100755
--- a/binary_search_tool/test/binary_search_tool_tester.py
+++ b/binary_search_tool/test/binary_search_tool_tester.py
@@ -107,6 +107,22 @@ class BisectingUtilsTest(unittest.TestCase):
     self.assertTrue(os.readlink(state_file) != first_state)
     self.assertFalse(os.path.exists(first_state))
 
+  def test_tmp_cleanup(self):
+    bss = binary_search_state.MockBinarySearchState(
+        get_initial_items='echo "0\n1\n2\n3"', switch_to_good='./switch_tmp.py',
+        file_args=True)
+    bss.SwitchToGood(['0', '1', '2', '3'])
+
+    tmp_file = None
+    with open('tmp_file', 'r') as f:
+      tmp_file = f.read()
+    os.remove('tmp_file')
+
+    self.assertFalse(os.path.exists(tmp_file))
+    ws = common.ReadWorkingSet()
+    for i in range(3):
+      self.assertEquals(ws[i], 42)
+
   def check_output(self):
     _, out, _ = command_executer.GetCommandExecuter().RunCommandWOutput(
         'tail -n 10 logs/binary_search_tool_tester.py.out')
@@ -132,6 +148,7 @@ def Main(argv):
   suite.addTest(BisectingUtilsTest('test_bad_install_script'))
   suite.addTest(BisectingUtilsTest('test_bad_save_state'))
   suite.addTest(BisectingUtilsTest('test_save_state'))
+  suite.addTest(BisectingUtilsTest('test_tmp_cleanup'))
   runner = unittest.TextTestRunner()
   runner.run(suite)
 
diff --git a/binary_search_tool/test/switch_tmp.py b/binary_search_tool/test/switch_tmp.py
new file mode 100755
index 00000000..165004ed
--- /dev/null
+++ b/binary_search_tool/test/switch_tmp.py
@@ -0,0 +1,34 @@
+#!/usr/bin/python2
+"""Change portions of the object files to good.
+
+This file is a test switch script. Used only for the test test_tmp_cleanup.
+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 42).
+"""
+
+from __future__ import print_function
+
+import sys
+
+import common
+
+
+def Main(argv):
+  working_set = common.ReadWorkingSet()
+  object_index = common.ReadObjectIndex(argv[1])
+
+  # Random number so the results can be checked
+  for oi in object_index:
+    working_set[int(oi)] = 42
+
+  common.WriteWorkingSet(working_set)
+  with open('tmp_file', 'w') as f:
+    f.write(argv[1])
+
+  return 0
+
+
+if __name__ == '__main__':
+  retval = Main(sys.argv)
+  sys.exit(retval)