From 9eabadd22d019cc02de771c6b77640d5746e0a0c Mon Sep 17 00:00:00 2001
From: David Rusu <davidrusu.me@gmail.com>
Date: Sat, 2 Nov 2024 00:01:20 +0400
Subject: [PATCH] cryptarchia: cleanup

---
 cryptarchia/cryptarchia.py      | 18 +++++++++---------
 cryptarchia/test_common.py      | 10 +++++++++-
 cryptarchia/test_fork_choice.py | 14 +++++++-------
 3 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/cryptarchia/cryptarchia.py b/cryptarchia/cryptarchia.py
index 3446f43..6a9feda 100644
--- a/cryptarchia/cryptarchia.py
+++ b/cryptarchia/cryptarchia.py
@@ -175,11 +175,11 @@ class Coin:
 
 @dataclass
 class MockLeaderProof:
-    commitment: Id = bytes(32)
-    nullifier: Id = bytes(32)
-    evolved_commitment: Id = bytes(32)
-    slot: Slot = field(default_factory=lambda: Slot(0))
-    parent: Id = bytes(32)
+    commitment: Id
+    nullifier: Id
+    evolved_commitment: Id
+    slot: Slot
+    parent: Id
 
     @staticmethod
     def new(coin: Coin, slot: Slot, parent: Id):
@@ -207,10 +207,10 @@ class MockLeaderProof:
 @dataclass
 class BlockHeader:
     slot: Slot
-    parent: Id = bytes(32)
-    content_size: int = 0
-    content_id: Id = bytes(32)
-    leader_proof: MockLeaderProof = field(default_factory=MockLeaderProof)
+    parent: Id
+    content_size: int
+    content_id: Id
+    leader_proof: MockLeaderProof
 
     orphaned_proofs: List["BlockHeader"] = field(default_factory=list)
 
diff --git a/cryptarchia/test_common.py b/cryptarchia/test_common.py
index 78cbd20..cd35090 100644
--- a/cryptarchia/test_common.py
+++ b/cryptarchia/test_common.py
@@ -50,7 +50,15 @@ def mk_config(initial_stake_distribution: list[Coin]) -> Config:
 
 def mk_genesis_state(initial_stake_distribution: list[Coin]) -> LedgerState:
     return LedgerState(
-        block=BlockHeader(slot=Slot(0), parent=bytes(32)),
+        block=BlockHeader(
+            slot=Slot(0),
+            parent=bytes(32),
+            content_size=0,
+            content_id=bytes(32),
+            leader_proof=MockLeaderProof.new(
+                Coin(sk=0, value=0), Slot(0), parent=bytes(32)
+            ),
+        ),
         nonce=bytes(32),
         commitments_spend={c.commitment() for c in initial_stake_distribution},
         commitments_lead={c.commitment() for c in initial_stake_distribution},
diff --git a/cryptarchia/test_fork_choice.py b/cryptarchia/test_fork_choice.py
index 2acaae7..4c1b39b 100644
--- a/cryptarchia/test_fork_choice.py
+++ b/cryptarchia/test_fork_choice.py
@@ -4,7 +4,6 @@ from copy import deepcopy
 from cryptarchia.cryptarchia import (
     block_weight,
     ghost_maxvalid_bg,
-    BlockHeader,
     Slot,
     Coin,
     Follower,
@@ -33,7 +32,7 @@ class TestForkChoice(TestCase):
 
         coin = Coin(sk=1, value=100)
 
-        b0 = BlockHeader(slot=Slot(0), parent=bytes(32))
+        b0 = mk_genesis_state([]).block
 
         b1A = mk_block(b0, 1, coin, content=b"b1A")
         b2A = mk_block(b1A, 2, coin, content=b"b2A")
@@ -105,7 +104,7 @@ class TestForkChoice(TestCase):
 
         coin = Coin(sk=1, value=100)
 
-        b0 = BlockHeader(slot=Slot(0), parent=bytes(32))
+        b0 = mk_genesis_state([]).block
 
         b1A = mk_block(b0, 1, coin, content=b"b1A")
         b2A = mk_block(b1A, 2, coin, content=b"b2A")
@@ -186,7 +185,7 @@ class TestForkChoice(TestCase):
 
         coin = Coin(sk=1, value=100)
 
-        b0 = BlockHeader(slot=Slot(0), parent=bytes(32))
+        b0 = mk_genesis_state([]).block
         b1 = mk_block(b0, 1, coin)
         b2 = mk_block(b1, 2, coin)
         b3 = mk_block(b2, 3, coin)
@@ -223,7 +222,7 @@ class TestForkChoice(TestCase):
 
         coin = Coin(sk=1, value=100)
 
-        b0 = BlockHeader(slot=Slot(0), parent=bytes(32))
+        b0 = mk_genesis_state([]).block
         b1 = mk_block(b0, 1, coin)
         b2 = mk_block(b1, 2, coin)
         b3 = mk_block(b2, 3, coin)
@@ -261,7 +260,8 @@ class TestForkChoice(TestCase):
 
     def test_fork_choice_long_sparse_chain(self):
         # The longest chain is not dense after the fork
-        genesis = BlockHeader(slot=Slot(0), parent=bytes(32))
+        genesis = mk_genesis_state([]).block
+
         short_coin, long_coin = Coin(sk=0, value=100), Coin(sk=1, value=100)
         common, long_coin = mk_chain(parent=genesis, coin=long_coin, slots=range(50))
 
@@ -304,7 +304,7 @@ class TestForkChoice(TestCase):
         # The longest chain is also the densest after the fork
         short_coin, long_coin = Coin(sk=0, value=100), Coin(sk=1, value=100)
         common, long_coin = mk_chain(
-            parent=BlockHeader(slot=Slot(0), parent=bytes(32)),
+            parent=mk_genesis_state([]).block,
             coin=long_coin,
             slots=range(1, 50),
         )