erc20.py

erc20.py

import random
from recordclass import recordclass

ADDRESSES = list(range(5))
MAX_PENDING = 5

Transfer = recordclass("Transfer", ["sender", "dst", "value"])

TransferFrom = recordclass("TransferFrom", ["sender", "src", "dst", "value"])

Approve = recordclass("Approve", ["sender", "spender", "value"])


class Model:
    def __init__(self):

        self.balance = {i: 5 for i in ADDRESSES}
        self.allowance = {
            (sender, spender): 0 for sender in ADDRESSES for spender in ADDRESSES
        }
        self.pending = []

    def submit(self, t):
        # print(f"submit {t}")
        self.pending.append(t)

    def commit(self, i):
        t = self.pending.pop(i)
        # print(f"commit {t}")
        if isinstance(t, Transfer):
            if t.value <= self.balance[t.sender]:
                self.balance[t.sender] -= t.value
                self.balance[t.dst] += t.value
        elif isinstance(t, TransferFrom):
            if (
                t.value <= self.balance[t.src]
                and t.value <= self.allowance[(t.src, t.dst)]
            ):
                self.balance[t.src] -= t.value
                self.balance[t.dst] += t.value
                self.allowance[(t.src, t.sender)] -= t.value

                def is_lesser_approve(e):
                    return (
                        isinstance(e, Approve)
                        and e.sender == t.src
                        and e.spender == t.sender
                        and e.value < t.value
                        and 0 < e.value
                    )

                if any(is_lesser_approve(e) for e in self.pending):
                    assert False, f"Found the event, {self.pending=}"
        elif isinstance(t, Approve):
            self.allowance[(t.sender, t.spender)] = t.value


def test_erc20():
    NUM_RUNS = 10000000
    K = 5

    count = 0
    for _ in range(NUM_RUNS):
        count += 1
        try:
            m = Model()
            # print("~~~")
            for _ in range(K):
                pending = len(m.pending)
                assert(pending <= 1)
                if MAX_PENDING <= pending or (1 <= pending and random.random() < 0.25):
                    i = random.randint(0, len(m.pending) - 1)
                    m.commit(i)
                else:
                    ctor = random.choice((Transfer, TransferFrom, Approve))
                    if ctor == Transfer:
                        sender = random.choice(ADDRESSES)
                        dst = random.choice(list(set(ADDRESSES) - {sender}))
                        value = random.choice(range(1, max(1, m.balance[sender]) + 1))
                        m.submit(ctor(sender, dst, value))
                    if ctor == TransferFrom:
                        sender = random.choice(ADDRESSES)
                        src = random.choice(list(set(ADDRESSES) - {sender}))
                        dst = random.choice(list(set(ADDRESSES) - {src}))
                        value = random.choice(range(1, max(1, m.balance[src]) + 1))
                        m.submit(ctor(sender, src, dst, value))
                    if ctor == Approve:
                        sender = random.choice(ADDRESSES)
                        spender = random.choice(list(set(ADDRESSES) - {sender}))
                        value = random.choice(range(0, 5))
                        m.submit(ctor(sender, spender, value))
        except AssertionError as e:
            print(f"I simulated {count} runs until one failed")
            raise e
            

if __name__ == "__main__":
    test_erc20()