1 files changed, 10 insertions, 34 deletions

diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index d53c613a3..3109f2b47 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -10,7 +10,7 @@ mod firmware_updater;
 mod mem_flash;
 mod partition;
 
-pub use boot_loader::{BootError, BootFlash, BootLoader, Flash, FlashConfig, MultiFlashConfig, SingleFlashConfig};
+pub use boot_loader::{BootError, BootFlash, BootLoader, FlashConfig, MultiFlashConfig, SingleFlashConfig};
 pub use firmware_updater::{FirmwareUpdater, FirmwareUpdaterError};
 pub use partition::Partition;
 
@@ -77,12 +77,8 @@ mod tests {
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
 
-        let mut magic = [0; 4];
         let mut page = [0; 4096];
-        assert_eq!(
-            State::Boot,
-            bootloader.prepare_boot(&mut flash, &mut magic, &mut page).unwrap()
-        );
+        assert_eq!(State::Boot, bootloader.prepare_boot(&mut flash, &mut page).unwrap());
     }
 
     #[test]
@@ -103,19 +99,14 @@ mod tests {
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
         let mut updater = FirmwareUpdater::new(DFU, STATE);
-        let mut offset = 0;
-        for chunk in update.chunks(4096) {
-            block_on(updater.write_firmware(offset, chunk, &mut flash)).unwrap();
-            offset += chunk.len();
-        }
+        block_on(updater.write_firmware(0, &update, &mut flash)).unwrap();
         block_on(updater.mark_updated(&mut flash, &mut aligned)).unwrap();
 
-        let mut magic = [0; 4];
-        let mut page = [0; 4096];
+        let mut page = [0; 1024];
         assert_eq!(
             State::Swap,
             bootloader
-                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut magic, &mut page)
+                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut page)
                 .unwrap()
         );
 
@@ -132,7 +123,7 @@ mod tests {
         assert_eq!(
             State::Swap,
             bootloader
-                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut magic, &mut page)
+                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut page)
                 .unwrap()
         );
 
@@ -150,7 +141,7 @@ mod tests {
         assert_eq!(
             State::Boot,
             bootloader
-                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut magic, &mut page)
+                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut page)
                 .unwrap()
         );
     }
@@ -176,25 +167,16 @@ mod tests {
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
 
-        let mut offset = 0;
-        for chunk in update.chunks(2048) {
-            block_on(updater.write_firmware(offset, chunk, &mut dfu)).unwrap();
-            offset += chunk.len();
-        }
+        block_on(updater.write_firmware(0, &update, &mut dfu)).unwrap();
         block_on(updater.mark_updated(&mut state, &mut aligned)).unwrap();
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
-        let mut magic = [0; 4];
         let mut page = [0; 4096];
 
         assert_eq!(
             State::Swap,
             bootloader
-                .prepare_boot(
-                    &mut MultiFlashConfig::new(&mut active, &mut state, &mut dfu),
-                    &mut magic,
-                    &mut page
-                )
+                .prepare_boot(&mut MultiFlashConfig::new(&mut active, &mut state, &mut dfu), &mut page)
                 .unwrap()
         );
 
@@ -229,22 +211,16 @@ mod tests {
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
 
-        let mut offset = 0;
-        for chunk in update.chunks(4096) {
-            block_on(updater.write_firmware(offset, chunk, &mut dfu)).unwrap();
-            offset += chunk.len();
-        }
+        block_on(updater.write_firmware(0, &update, &mut dfu)).unwrap();
         block_on(updater.mark_updated(&mut state, &mut aligned)).unwrap();
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
-        let mut magic = [0; 4];
         let mut page = [0; 4096];
         assert_eq!(
             State::Swap,
             bootloader
                 .prepare_boot(
                     &mut MultiFlashConfig::new(&mut active, &mut state, &mut dfu,),
-                    &mut magic,
                     &mut page
                 )
                 .unwrap()