Add tests

tests/native/test_state.py

@@ -1,6 +1,7 @@
 import unittest
 import os
 
+from manticore.core.state import StateBase
 from manticore.utils.event import Eventful
 from manticore.platforms import linux
 from manticore.native.state import State
@@ -157,6 +158,44 @@ def test_tainted_symbolic_value(self):
         expr = self.state.new_symbolic_value(64, taint=taint)
         self.assertEqual(expr.taint, frozenset(taint))
 
+    def test_state_hook(self):
+        initial_state = State(ConstraintSet(), FakePlatform())
+
+        def fake_hook(_: StateBase) -> None:
+            return None
+
+        self.assertTrue(len(initial_state._hooks) == 0)
+        self.assertTrue(len(initial_state._after_hooks) == 0)
+
+        with initial_state as new_state:
+            self.assertTrue(len(new_state._hooks) == 0)
+            self.assertTrue(len(new_state._after_hooks) == 0)
+
+            new_state.add_hook(0x4000, fake_hook, after=False)
+
+            new_state.add_hook(0x4001, fake_hook, after=True)
+            new_state.add_hook(0x4001, fake_hook, after=False)
+
+            self.assertTrue(len(new_state._hooks) == 2)
+            self.assertTrue(len(new_state._after_hooks) == 1)
+            self.assertTrue(len(initial_state._hooks) == 0)
+            self.assertTrue(len(initial_state._after_hooks) == 0)
+
+            # Remove one of the hooks we added
+            new_state.remove_hook(0x4000, fake_hook, after=False)
+            # Try to remove a non-existent hook
+            self.assertFalse(new_state.remove_hook(0x4000, fake_hook, after=True))
+
+            self.assertTrue(len(new_state._hooks) == 1)
+            self.assertTrue(len(new_state._after_hooks) == 1)
+            self.assertTrue(len(initial_state._hooks) == 0)
+            self.assertTrue(len(initial_state._after_hooks) == 0)
+
+            # Add to all PC
+            initial_state.add_hook(None, fake_hook, after=True)
+
+        self.assertTrue(len(initial_state._after_hooks) == 1)
+
     def testContextSerialization(self):
         import pickle as pickle
 
@@ -211,6 +250,31 @@ def testContextSerialization(self):
         self.assertEqual(new_new_state.context["step"], 30)
 
 
+class StateHooks(unittest.TestCase):
+    def setUp(self):
+        core = config.get_group("core")
+        core.seed = 61
+        core.mprocessing = core.mprocessing.single
+
+        dirname = os.path.dirname(__file__)
+        self.m = Manticore(os.path.join(dirname, "binaries", "basic_linux_amd64"), policy="random")
+
+    def test_state_hooks(self):
+        def do_nothing(_: StateBase) -> None:
+            return None
+
+        def fin(_: StateBase) -> None:
+            return None
+
+        @self.m.hook(0x400610, after=True)
+        def process_hook(state: State) -> None:
+            state.add_hook(None, do_nothing, after=False)
+            state.add_hook(0x400647, fin, after=True)
+            self.assertTrue(state.remove_hook(0x400610, process_hook, after=True))
+
+        self.m.run()
+
+
 class StateMergeTest(unittest.TestCase):
 
     # Need to add a plugin that counts the number of states in did_fork_state, and records the max