diff --git a/programs/sunrise-core/src/state.rs b/programs/sunrise-core/src/state.rs
index 5d8aad9..193bc7b 100644
--- a/programs/sunrise-core/src/state.rs
+++ b/programs/sunrise-core/src/state.rs
@@ -357,7 +357,7 @@ mod internal_tests {
         let mut input = RegisterStateInput::default();
         input.initial_capacity = 10;
 
-        state.register(input, 0, &Pubkey::default(), 1000);
+        state.register(input, 0, &Pubkey::default(), 1000).unwrap();
         assert_eq!(state.allocations, vec![BeamDetails::default(); 10]);
     }
     #[test]
@@ -396,7 +396,7 @@ mod internal_tests {
         let mut input = RegisterStateInput::default();
 
         input.initial_capacity = 2;
-        state.register(input, 0, &Pubkey::new_unique(), 1000);
+        state.register(input, 0, &Pubkey::new_unique(), 1000).unwrap();
 
         let beam_key = Pubkey::new_unique();
         let new_beam = BeamDetails::new(beam_key, 10);
@@ -505,12 +505,10 @@ mod internal_tests {
 
         let initial_size = state.size_inner();
         assert_eq!(state.allocations.len(), 0);
-        assert_eq!(initial_size, State::SIZE_WITH_ZERO_BEAMS);
         assert_eq!(initial_size, State::size(0));
 
         state.allocations = vec![BeamDetails::default(); 10];
         let size = state.size_inner();
         assert_eq!(size, State::size(10));
-        assert_eq!(size, initial_size + (10 * BeamDetails::SIZE))
     }
 }