angr-angr/tests/misc/test_hookset.py

#!/usr/bin/env python3
# pylint: disable=missing-class-docstring,disable=no-self-use
from __future__ import annotations
import unittest

import angr


class TestHookSet(unittest.TestCase):
    def test_hookset(self):
        class Foo:
            def run(self):
                return self.blah()

            def blah(self):  # pylint:disable=no-self-use
                return ["foo"]

            def install_hooks(self, tech):
                angr.misc.HookSet.install_hooks(self, blah=tech.blah)

            def remove_hooks(self, tech):
                angr.misc.HookSet.remove_hooks(self, blah=tech.blah)

        class Bar:
            def blah(self, foo):  # pylint:disable=no-self-use
                return ["bar", *foo.blah()]

        class Baz:
            def blah(self, foo):  # pylint:disable=no-self-use
                return ["baz", *foo.blah()]

        class Coward:
            def blah(self, foo):  # pylint:disable=no-self-use,unused-argument
                return ["coward"]

        foo = Foo()
        assert foo.run() == ["foo"]

        bar = Bar()
        baz = Baz()
        foo.install_hooks(bar)
        foo.install_hooks(baz)
        assert foo.run() == ["baz", "bar", "foo"]

        foo.remove_hooks(bar)
        foo.remove_hooks(baz)
        assert foo.run() == ["foo"]

        coward = Coward()
        foo.install_hooks(coward)
        assert foo.run() == ["coward"]

        foo.remove_hooks(coward)
        assert foo.run() == ["foo"]


if __name__ == "__main__":
    unittest.main()