onionrblocks/tests/test_anonvdf.py

import unittest
from onionrblocks.generators import anonvdf
import kasten
import mimcvdf

from time import time
from math import floor

def get_rounds_for_ttl_seconds(seconds: int, size_bytes: int):
    second_cost = 4
    byte_cost = 10
    return (seconds * second_cost) + (size_bytes * byte_cost)

class TestAnonVDF(unittest.TestCase):

    def test_vdf_create(self):
        ttl = 10000
        test_data = kasten.generator.pack.pack(b"test", "tst",  app_metadata={"ttl": ttl})
        test_vdf = mimcvdf.vdf_create(test_data, get_rounds_for_ttl_seconds(ttl, len(test_data)), dec=True)

        generated = anonvdf.AnonVDFGenerator.generate(test_data, get_rounds_for_ttl_seconds(ttl, len(test_data)))
        self.assertEqual(
            generated.get_packed(),
            test_data
        )
        self.assertEqual(generated.id, int(test_vdf).to_bytes(64, byteorder="big"))

    def test_vdf_block_validate_not_enough(self):

        t = floor(time()) - 60
        ttl = 3600
        test_data = kasten.generator.pack.pack(b"test", "tst", app_metadata={"ttl": ttl}, timestamp=t)
        rds = get_rounds_for_ttl_seconds(3500, len(test_data))
        test_id = mimcvdf.vdf_create(test_data, rds)
        self.assertRaises(anonvdf.InvalidID, anonvdf.AnonVDFGenerator.validate_id, test_id, test_data)

    def test_vdf_incermental(self):
        last = 1
        ttl = 30000
        for i in range(1, 5):
            t = time()
            rds = 0
            test_data = kasten.generator.pack.pack(b"test" * i, "tst", app_metadata={"ttl": ttl})
            test_id = mimcvdf.vdf_create(test_data, get_rounds_for_ttl_seconds(ttl, len(test_data)))
            newT = time()
            self.assertGreater(newT - t, last)
            print(newT - t)
            anonvdf.AnonVDFGenerator.validate_id(test_id, test_data)

    def test_block_invalid_ttl(self):
        return

    def test_vdf_rounds_seconds(self):
        size = 10000
        per_byte = 10
        per_second = 4
        expected_rounds = (3600 * per_second) + (size * per_byte)
        self.assertEqual(get_rounds_for_ttl_seconds(3600, size), expected_rounds)


unittest.main()