diff --git a/lockdrop/calculations.py b/lockdrop/calculations.py
index fdfec26..19afb67 100644
--- a/lockdrop/calculations.py
+++ b/lockdrop/calculations.py
@@ -179,7 +179,7 @@ def generate_test_output(final_data):
             f"{year}_years": int(final_galaxy_allocations[year] * Decimal('1e8'))
             for year in range(1, 6)
         },
-        "total_allocation_sz": int(total_allocation * Decimal('1e8'))
+        "total": int(total_allocation * Decimal('1e8'))
     }
 
     return output
diff --git a/tests/base_test.py b/tests/base_test.py
index 15c239a..22eab8d 100644
--- a/tests/base_test.py
+++ b/tests/base_test.py
@@ -31,7 +31,8 @@ class BaseAllocationTest(unittest.TestCase):
         with open(f'{cls.generated_dir}/generated-participants.json', 'r') as f:
             cls.participants = json.load(f)
         with open('lockdrop_allocations_notebook.json', 'r') as f:
-            cls.notebook_allocations = json.load(f)
+            notebook_data = json.load(f)
+            cls.notebook_allocations = notebook_data['allocations']
 
         cls.points_by_duration = cls._get_first_points()