function call instrumentation (for testing)

mainly used to count calls to Encrypted.encrypt()

rainbowadn/encryption/encrypted.py

@@ -3,9 +3,9 @@ from typing import Generic, Iterable, TypeVar
 from nacl.bindings import crypto_hash_sha256
 from nacl.secret import SecretBox
 
-from rainbowadn.hashing.mentionable import Mentionable
 from rainbowadn.hashing.hashpoint import HashPoint
 from rainbowadn.hashing.hashresolver import HashResolver
+from rainbowadn.hashing.mentionable import Mentionable
 from rainbowadn.hashing.origin import Origin
 from rainbowadn.hashing.rainbow_factory import RainbowFactory
 from rainbowadn.hashing.recursivementionable import RecursiveMentionable
@@ -42,11 +42,8 @@ class Encrypted(RecursiveMentionable, Generic[EncryptedType]):
     def points(self) -> Iterable[HashPoint]:
         return self.resolution
 
-    ecc = 0
-
     @classmethod
     def encrypt(cls, decrypted: EncryptedType, key: bytes) -> 'Encrypted[EncryptedType]':
-        cls.ecc += 1
         assert isinstance(key, bytes)
         hashpoints = tuple(decrypted.points()) if isinstance(decrypted, RecursiveMentionable) else ()
         resolution = tuple(

rainbowadn/testing/instrumentation.py

@@ -0,0 +1,34 @@
+import functools
+
+
+class Instrumentation:
+    def __init__(self, target, methodname: str):
+        self.target = target
+        self.methodname = methodname
+
+    def instrument(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def __enter__(self):
+        assert not hasattr(self, 'method')
+        self.method = getattr(self.target, self.methodname)
+
+        @functools.wraps(self.method)
+        def wrap(*args, **kwargs):
+            self.instrument(*args, **kwargs)
+            return self.method(*args, **kwargs)
+
+        setattr(self.target, self.methodname, wrap)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        setattr(self.target, self.methodname, self.method)
+        del self.method
+
+
+class Counter(Instrumentation):
+    def __init__(self, target, methodname: str):
+        super().__init__(target, methodname)
+        self.counter = 0
+
+    def instrument(self, *args, **kwargs):
+        self.counter += 1

rainbowadn/testing/test_all.py

@@ -22,6 +22,7 @@ from rainbowadn.hashing.nullability.notnull import NotNull
 from rainbowadn.hashing.rainbow_factory import RainbowFactory
 from rainbowadn.hashing.resolvermetaorigin import ResolverMetaOrigin
 from rainbowadn.testing.dictresolver import DictResolver
+from rainbowadn.testing.instrumentation import Counter
 from rainbowadn.v13.algo import MINT_CONST
 from rainbowadn.v13.bankchain import BankChain
 from rainbowadn.v13.subject import Subject
@@ -156,6 +157,7 @@ class TestAll(unittest.TestCase):
         print(tree.loose().reference.reference.resolve().resolve().key.resolve().metadata.resolve().integer)
 
     def test_encryption(self):
+        instrumentation = Counter(Encrypted, 'encrypt')
         with self.subTest('setup'):
             key = b'a' * 32
             dr = DictResolver()
@@ -169,8 +171,9 @@ class TestAll(unittest.TestCase):
             print(tree.reference.str(0))
         with self.subTest('encrypt'):
             target = tree.reference
-            target = Encrypted.encrypt(target, key).decrypted
-            print(Encrypted.ecc)
+            with instrumentation:
+                target = Encrypted.encrypt(target, key).decrypted
+                print(instrumentation.counter)
             tree = tree.create(target)
             print(tree.reference.str(0))
         with self.subTest('alter'):
@@ -179,12 +182,14 @@ class TestAll(unittest.TestCase):
             print(tree.reference.str(0))
         with self.subTest('encrypt and migrate'):
             target = tree.reference
-            eeed = Encrypted.encrypt(target, key)
-            print(Encrypted.ecc)
+            with instrumentation:
+                eeed = Encrypted.encrypt(target, key)
+                print(instrumentation.counter)
             dr.save(HashPoint.of(eeed))
             print(ResolverMetaOrigin(dr).migrate(HashPoint.of(eeed)).resolve().decrypted.str(0))
         with self.subTest('re-encrypt'):
             new_key = b'b' * 32
             target = eeed.decrypted
-            Encrypted.encrypt(target, new_key)
-            print(Encrypted.ecc)
+            with instrumentation:
+                Encrypted.encrypt(target, new_key)
+                print(instrumentation.counter)