Add unit tests and update binaryview save

1 files changed, 76 insertions, 0 deletions

diff --git a/suite/testcommon.py b/suite/testcommon.py
index 06ef1b5b..628e129f 100644
--- a/suite/testcommon.py
+++ b/suite/testcommon.py
@@ -926,6 +926,82 @@ class VerifyBuilder(Builder):
         finally:
             self.delete_package("helloworld")
 
+    def test_verify_persistent_undo(self):
+        file_name = self.unpackage_file("helloworld")
+        try:
+            temp_name = next(tempfile._get_candidate_names()) + ".bndb"
+
+            bv = binja.BinaryViewType['ELF'].open(file_name)
+            bv.update_analysis_and_wait()
+
+            bv.begin_undo_actions()
+            bv.functions[0].set_comment(bv.functions[0].start, "Function start")
+            bv.functions[0].set_comment(bv.functions[0].start, "Function start!")
+            bv.commit_undo_actions()
+
+            bv.begin_undo_actions()
+            bv.add_function(bv.functions[0].start + 4)
+            bv.commit_undo_actions()
+
+            comments = self.get_comments(bv)
+            functions = self.get_functions(bv)
+
+            bv.create_database(temp_name)
+            bv.file.close()
+            del bv
+
+            bv = binja.FileMetadata(temp_name).open_existing_database(temp_name).get_view_of_type('ELF')
+            bv.update_analysis_and_wait()
+
+            bv.undo()
+            bv.undo()
+
+            bndb_functions = self.get_functions(bv)
+            bndb_comments = self.get_comments(bv)
+
+            bv.file.close()
+            del bv
+            os.unlink(temp_name)
+
+            return functions == bndb_functions and comments == bndb_comments
+        finally:
+            self.delete_package("helloworld")
+
+    def test_verify_clean_save(self):
+        file_name = self.unpackage_file("helloworld")
+        try:
+            temp_name = next(tempfile._get_candidate_names()) + ".bndb"
+
+            bv = binja.BinaryViewType['ELF'].open(file_name)
+            bv.update_analysis_and_wait()
+
+            bv.begin_undo_actions()
+            bv.functions[0].set_comment(bv.functions[0].start, "This is a secret comment")
+            bv.commit_undo_actions()
+
+            bv.begin_undo_actions()
+            bv.functions[0].set_comment(bv.functions[0].start, "Function start!")
+            bv.commit_undo_actions()
+
+            bv.create_database(temp_name, clean=True)
+            bv.file.close()
+            del bv
+
+            bv = binja.FileMetadata(temp_name).open_existing_database(temp_name).get_view_of_type('ELF')
+            bv.update_analysis_and_wait()
+
+            bv.undo()
+
+            comment = bv.functions[0].get_comment_at(bv.functions[0].start)
+
+            bv.file.close()
+            del bv
+            os.unlink(temp_name)
+
+            return comment == "Function start!"
+        finally:
+            self.delete_package("helloworld")
+
     def test_memory_leaks(self):
         """Detected memory leaks during analysis"""
         # This test will attempt to detect object leaks during headless analysis