2 files changed, 48 insertions, 1 deletions

diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index a2259411f..4c28d7aa4 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -313,7 +313,8 @@ mod tests {
         ))
         .is_ok());
     }
-    struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize>([u8; SIZE]);
+
+    pub struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize>(pub [u8; SIZE]);
 
     impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFlash
         for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
diff --git a/embassy-boot/boot/src/partition.rs b/embassy-boot/boot/src/partition.rs
index 9918fb836..3ccd4dd76 100644
--- a/embassy-boot/boot/src/partition.rs
+++ b/embassy-boot/boot/src/partition.rs
@@ -102,3 +102,49 @@ impl Partition {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::MemFlash;
+    use crate::Partition;
+
+    #[test]
+    fn can_erase() {
+        let mut flash = MemFlash::<1024, 64, 4>([0x00; 1024]);
+        let partition = Partition::new(256, 512);
+
+        partition.erase_blocking(&mut flash, 64, 192).unwrap();
+
+        for (index, byte) in flash.0.iter().copied().enumerate().take(256 + 64) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.0.iter().copied().enumerate().skip(256 + 64).take(128) {
+            assert_eq!(0xFF, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.0.iter().copied().enumerate().skip(256 + 64 + 128) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+    }
+
+    #[test]
+    fn can_wipe() {
+        let mut flash = MemFlash::<1024, 64, 4>([0x00; 1024]);
+        let partition = Partition::new(256, 512);
+
+        partition.wipe_blocking(&mut flash).unwrap();
+
+        for (index, byte) in flash.0.iter().copied().enumerate().take(256) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.0.iter().copied().enumerate().skip(256).take(256) {
+            assert_eq!(0xFF, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.0.iter().copied().enumerate().skip(512) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+    }
+}