lockdrop-simulation/tests/test_allocations.py
2025-08-04 14:10:17 +05:30

141 lines
5.6 KiB
Python

import unittest
from tabulate import tabulate
from base_test import BaseAllocationTest
class AllocationTest(BaseAllocationTest):
"""Test allocation comparisons between notebook and zenithd"""
def test_0_star_allocations(self):
"""Test star allocations for all lock periods"""
print("\nSTAR ALLOCATIONS COMPARISON")
table_data = []
headers = ["Lock Period", "Point", "Zenith Address", "Notebook ($sZ)", "zenithd ($sZ)", "Difference"]
for lock_period in sorted(self.points_by_duration.keys()):
with self.subTest(lock_period=lock_period):
duration_data = self.points_by_duration[lock_period]
star_data = duration_data['star']
if star_data == None:
continue
zenith_addr = star_data['zenith_address']
self.assertIsNotNone(zenith_addr,
f"Could not find zenith address for star {star_data['point']}")
allocation_data = self._get_point_allocation_from_api(zenith_addr, star_data['point'])
self.assertIsNotNone(allocation_data,
f"zenithd returned no allocation for star {star_data['point']}")
api_allocation = int(allocation_data['allocated_amount']['amount'])
notebook_key = f'{lock_period}_years'
notebook_allocation = self.notebook_allocations['stars'].get(notebook_key, 0)
difference = api_allocation - notebook_allocation
table_data.append([
f"{lock_period} years",
star_data['point'],
zenith_addr,
f"{notebook_allocation:,}",
f"{api_allocation:,}",
f"{difference:+,}" if difference != 0 else "0"
])
self.assertEqual(notebook_allocation, api_allocation,
f"Star {star_data['point']} ({lock_period}Y): "
f"Notebook={notebook_allocation:,} $sZ, "
f"zenithd={api_allocation:,} $sZ, "
f"Diff={difference:+,} $sZ")
print(tabulate(table_data, headers=headers, tablefmt="grid"))
def test_1_galaxy_allocations(self):
"""Test galaxy allocations for all lock periods"""
print("\nGALAXY ALLOCATIONS COMPARISON")
table_data = []
headers = ["Lock Period", "Point", "Zenith Address", "Notebook ($sZ)", "zenithd ($sZ)", "Difference"]
for lock_period in sorted(self.points_by_duration.keys()):
with self.subTest(lock_period=lock_period):
duration_data = self.points_by_duration[lock_period]
galaxy_data = duration_data['galaxy']
if galaxy_data == None:
continue
zenith_addr = galaxy_data['zenith_address']
self.assertIsNotNone(zenith_addr,
f"Could not find zenith address for galaxy {galaxy_data['point']}")
allocation_data = self._get_point_allocation_from_api(zenith_addr, galaxy_data['point'])
self.assertIsNotNone(allocation_data,
f"zenithd returned no allocation for galaxy {galaxy_data['point']}")
api_allocation = int(allocation_data['allocated_amount']['amount'])
notebook_key = f'{lock_period}_years'
notebook_allocation = self.notebook_allocations['galaxies'].get(notebook_key, 0)
difference = api_allocation - notebook_allocation
table_data.append([
f"{lock_period} years",
galaxy_data['point'],
zenith_addr,
f"{notebook_allocation:,}",
f"{api_allocation:,}",
f"{difference:+,}" if difference != 0 else "0"
])
self.assertEqual(notebook_allocation, api_allocation,
f"Galaxy {galaxy_data['point']} ({lock_period}Y): "
f"Notebook={notebook_allocation:,} $sZ, "
f"zenithd={api_allocation:,} $sZ, "
f"Diff={difference:+,} $sZ")
print(tabulate(table_data, headers=headers, tablefmt="grid"))
def test_2_total_allocations(self):
"""Test total allocations for all participants"""
print("\nTOTAL ALLOCATIONS COMPARISON")
notebook_total = self.notebook_allocations.get('total', 0)
if notebook_total == 0:
self.skipTest("No total_allocation found in notebook data")
api_total = 0
for i, participant in enumerate(self.participants):
zenith_address = participant['attestation']['payload']['payload']['address']
with self.subTest(participant=i, address=zenith_address):
participant_total = self._get_total_address_allocation_amount_from_api(zenith_address)
api_total += participant_total
# Summary table
difference = api_total - notebook_total
table_data = [
["Notebook", f"{notebook_total:,}"],
["zenithd", f"{api_total:,}"],
["Difference", f"{difference:+,}" if difference != 0 else "0"]
]
headers = ["Source", "Total Allocation ($sZ)"]
print(tabulate(table_data, headers=headers, tablefmt="grid"))
self.assertEqual(notebook_total, api_total,
f"Total allocation mismatch: "
f"Notebook={notebook_total:,} $sZ, "
f"zenithd={api_total:,} $sZ, "
f"Diff={difference:+,} $sZ")
if __name__ == "__main__":
unittest.main(verbosity=2)