diff --git a/src/libwolfboot.c b/src/libwolfboot.c
index 01e0870646..d6e4109bf1 100644
--- a/src/libwolfboot.c
+++ b/src/libwolfboot.c
@@ -752,17 +752,27 @@ void RAMFUNCTION wolfBoot_update_trigger(void)
      * FLAGS_INVERT needs erased flags because the bin-assemble's fill byte may
      * not match what's in wolfBoot */
     if (FLAGS_UPDATE_EXT()) {
-        ext_flash_erase(lastSector, SECTOR_FLAGS_SIZE);
+        ext_flash_erase(lastSector, WOLFBOOT_SECTOR_SIZE);
+        wolfBoot_set_partition_state(PART_UPDATE, st);
     } else {
 #ifndef NVM_FLASH_WRITEONCE
-        hal_flash_erase(lastSector, SECTOR_FLAGS_SIZE);
+        hal_flash_erase(lastSector, WOLFBOOT_SECTOR_SIZE);
+        wolfBoot_set_partition_state(PART_UPDATE, st);
 #else
+        uint32_t magic = WOLFBOOT_MAGIC_TRAIL;
+        uint32_t offset = SECTOR_FLAGS_SIZE;
+#ifdef FLAGS_HOME
+        offset -= (PART_BOOT_ENDFLAGS - PART_UPDATE_ENDFLAGS);
+#endif
         selSec = nvm_select_fresh_sector(PART_UPDATE);
         XMEMCPY(NVM_CACHE, (uint8_t*)lastSector - WOLFBOOT_SECTOR_SIZE * selSec,
             WOLFBOOT_SECTOR_SIZE);
         /* write to the non selected sector */
         hal_flash_erase(lastSector - WOLFBOOT_SECTOR_SIZE * !selSec,
             WOLFBOOT_SECTOR_SIZE);
+
+        NVM_CACHE[offset] = IMG_STATE_UPDATING;
+        memcpy(NVM_CACHE + offset + 1, &magic, sizeof(uint32_t));
         hal_flash_write(lastSector - WOLFBOOT_SECTOR_SIZE * !selSec, NVM_CACHE,
             WOLFBOOT_SECTOR_SIZE);
         /* erase the previously selected sector */
@@ -771,7 +781,6 @@ void RAMFUNCTION wolfBoot_update_trigger(void)
 #endif
     }
 
-    wolfBoot_set_partition_state(PART_UPDATE, st);
 
     if (FLAGS_UPDATE_EXT()) {
         ext_flash_lock();
diff --git a/tools/unit-tests/unit-enc-nvm.c b/tools/unit-tests/unit-enc-nvm.c
index d5e361a6c6..4fec9b20ce 100644
--- a/tools/unit-tests/unit-enc-nvm.c
+++ b/tools/unit-tests/unit-enc-nvm.c
@@ -27,6 +27,7 @@
 #define MOCK_ADDRESS_SWAP 0xCE000000
 const char ENCRYPT_KEY[] = "0123456789abcdef0123456789abcdef0123456789abcdef";
 #include <stdio.h>
+#include "image.h"
 #include "encrypt.h"
 #include "libwolfboot.c"
 #include <fcntl.h>
@@ -274,10 +275,10 @@ START_TEST (test_nvm_update_with_encryption)
     /* Triggering update to set flags */
     wolfBoot_update_trigger();
 
-    /* Current selected should now be 0 */
+    /* Current selected should now be 1 */
     ret = nvm_select_fresh_sector(PART_UPDATE);
-    ck_assert_msg(ret == 0, "Failed to select updating fresh sector\n");
-    ck_assert_msg(erased_nvm_bank1 != 0, "Did not erase the non-selected bank");
+    ck_assert_msg(ret == 1, "Failed to select updating fresh sector\n");
+    ck_assert_msg(erased_nvm_bank0 != 0, "Did not erase the non-selected bank");
 
     magic = get_partition_magic(PART_UPDATE);
     ck_assert_msg(*magic == *boot_word,
diff --git a/tools/unit-tests/unit-nvm.c b/tools/unit-tests/unit-nvm.c
index 774dc64271..74f40144b8 100644
--- a/tools/unit-tests/unit-nvm.c
+++ b/tools/unit-tests/unit-nvm.c
@@ -269,10 +269,10 @@ START_TEST (test_nvm_select_fresh_sector)
     /* Triggering update to set flags */
     wolfBoot_update_trigger();
 
-    /* Current selected should now be 0 */
+    /* Current selected should now be 1 */
     ret = nvm_select_fresh_sector(PART_UPDATE);
-    ck_assert_msg(ret == 0, "Failed to select updating fresh sector\n");
-    ck_assert_msg(erased_nvm_bank1 > 0, "Did not erase the non-selected bank");
+    ck_assert_msg(ret == 1, "Failed to select updating fresh sector\n");
+    ck_assert_msg(erased_nvm_bank0 > 0, "Did not erase the non-selected bank");
 
     magic = get_partition_magic(PART_UPDATE);
     ck_assert_msg(*magic == *boot_word,