add tests to pot_iterations.rs

Chia-Network · Jan 24, 2025 · 31bdbc7 · 31bdbc7
1 parent 2bb5c1e
commit 31bdbc7
Showing 1 changed file with 101 additions and 0 deletions.
diff --git a/crates/chia-protocol/src/pot_iterations.rs b/crates/chia-protocol/src/pot_iterations.rs
@@ -69,6 +69,107 @@ pub fn calculate_ip_iters(
     )
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+    static NUM_SPS_SUB_SLOT: u32 = 32;
+    static NUM_SPS_SUB_SLOT_U8: u8 = 32;
+    static NUM_SP_INTERVALS_EXTRA: u8 = 3;
+
+    #[test]
+    fn test_is_overflow_block() {
+        assert!(!is_overflow_block(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, 27).expect("valid SP index"));
+        assert!(!is_overflow_block(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, 28).expect("valid SP index"));
+        assert!(is_overflow_block(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, 29).expect("valid SP index"));
+        assert!(is_overflow_block(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, 30).expect("valid SP index"));
+        assert!(is_overflow_block(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, 31).expect("valid SP index"));
+        assert!(matches!(
+            is_overflow_block(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, 32),
+            Err(_)
+        ));
+    }
+
+    #[test]
+    fn test_calculate_sp_iters() {
+        let ssi: u64 = 100001 * 64 * 4;
+        assert!(matches!(
+            calculate_sp_iters(NUM_SPS_SUB_SLOT, ssi, 32),
+            Err(_)
+        ));
+        calculate_sp_iters(NUM_SPS_SUB_SLOT, ssi, 31).expect("valid_result");
+    }
+
+    #[test]
+    fn test_calculate_ip_iters(){
+        // # num_sps_sub_slot: u32,
+        // # num_sp_intervals_extra: u8,
+        // # sub_slot_iters: u64,
+        // # signage_point_index: u8,
+        // # required_iters: u64,
+        let ssi: u64 = 100001 * 64 * 4;
+        let sp_interval_iters = ssi / NUM_SPS_SUB_SLOT as u64;
+
+        // Invalid signage point index
+        assert!(matches!(
+            calculate_ip_iters(
+                NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, 123, 100000
+            ),
+            Err(_)
+        ));
+
+        let sp_iters = sp_interval_iters * 13;
+
+        // required_iters too high
+        assert!(matches!(
+            calculate_ip_iters(
+                NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, sp_interval_iters.try_into().unwrap(), sp_interval_iters
+            ),
+            Err(_)
+        ));
+
+        // required_iters too high
+        assert!(matches!(
+            calculate_ip_iters(
+                NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, sp_interval_iters.try_into().unwrap(), sp_interval_iters * 12
+            ),
+            Err(_)
+        ));
+
+        // required_iters too low (0)
+        assert!(matches!(
+            calculate_ip_iters(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, sp_interval_iters.try_into().unwrap(), 0),
+            Err(_)
+        ));
+
+        let required_iters = sp_interval_iters - 1;
+        let ip_iters = calculate_ip_iters(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, 13, required_iters).expect("should be valid");
+        assert_eq!(ip_iters, sp_iters + (NUM_SP_INTERVALS_EXTRA as u64 * sp_interval_iters) + required_iters);
+
+        let required_iters = 1_u64;
+        let ip_iters = calculate_ip_iters(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, 13, required_iters).expect("valid");
+        assert_eq!(ip_iters, sp_iters + (NUM_SP_INTERVALS_EXTRA as u64 * sp_interval_iters) + required_iters);
+
+        let required_iters: u64 = ssi * 4 / 300;
+        let ip_iters = calculate_ip_iters(NUM_SPS_SUB_SLOT, NUM_SP_INTERVALS_EXTRA, ssi, 13, required_iters).expect("valid");
+        assert_eq!(ip_iters, sp_iters + (NUM_SP_INTERVALS_EXTRA as u64 * sp_interval_iters) + required_iters);
+        assert!(sp_iters < ip_iters);
+
+        // Overflow
+        let sp_iters = sp_interval_iters * (NUM_SPS_SUB_SLOT - 1) as u64;
+        let ip_iters = calculate_ip_iters(
+            NUM_SPS_SUB_SLOT, 
+            NUM_SP_INTERVALS_EXTRA,
+            ssi,
+            NUM_SPS_SUB_SLOT_U8 - 1_u8,
+            required_iters,
+        ).expect("valid");
+        assert_eq!(ip_iters, (sp_iters + (NUM_SP_INTERVALS_EXTRA as u64 * sp_interval_iters) + required_iters) % ssi);
+        assert!(sp_iters > ip_iters)
+    }
+
+
+}
+
 // TODO: enable and fix below
 
 // #[cfg(feature = "py-bindings")]