diff --git a/src/uart_flash.c b/src/uart_flash.c
index 3623c9c9..7e7b3b98 100644
--- a/src/uart_flash.c
+++ b/src/uart_flash.c
@@ -47,17 +47,22 @@
 int uart_tx(const uint8_t c);
 int uart_rx(uint8_t *c);
 
-static int wait_ack(void)
+static int wait_ack_cycles(int cycles)
 {
+    uint8_t c;
     volatile int count = 0;
-    while(++count < WAIT_CYCLES) {
-        uint8_t c;
+    while(++count < cycles) {
         if ((uart_rx(&c) == 1) && (c == CMD_ACK))
             return 0;
     }
     return -1;
 }
 
+static int wait_ack(void)
+{
+    return wait_ack_cycles(WAIT_CYCLES);
+}
+
 static int uart_rx_timeout(uint8_t *c)
 {
     volatile int count = 0;
@@ -144,7 +149,7 @@ int  ext_flash_erase(uintptr_t address, int len)
             return -1;
     }
     /* Wait for extra ack at the end of Erase */
-    if (wait_ack() == 0)
+    if (wait_ack_cycles(WAIT_CYCLES * ERASE_TIMEOUT) == 0)
         return 0;
     return -1;
 }
diff --git a/tools/unit-tests/unit-uart-flash.c b/tools/unit-tests/unit-uart-flash.c
index c26e11e1..51b75b1c 100644
--- a/tools/unit-tests/unit-uart-flash.c
+++ b/tools/unit-tests/unit-uart-flash.c
@@ -69,12 +69,45 @@ START_TEST(test_ext_flash_read_timeout_returns_error)
 }
 END_TEST
 
+START_TEST(test_ext_flash_erase_success)
+{
+    uint8_t script[11];
+    int ret;
+
+    /* 10 command ACKs + the erase-completion ACK */
+    memset(script, CMD_ACK, sizeof(script));
+    reset_uart_script(script, sizeof(script));
+
+    ret = ext_flash_erase(0x1000, 0x1000);
+
+    ck_assert_int_eq(ret, 0);
+    ck_assert_int_eq(rx_script_pos, 11);
+}
+END_TEST
+
+START_TEST(test_ext_flash_erase_timeout_returns_error)
+{
+    uint8_t script[10];
+    int ret;
+
+    /* Command ACKs only: the erase-completion ACK never arrives */
+    memset(script, CMD_ACK, sizeof(script));
+    reset_uart_script(script, sizeof(script));
+
+    ret = ext_flash_erase(0x1000, 0x1000);
+
+    ck_assert_int_eq(ret, -1);
+}
+END_TEST
+
 Suite *wolfboot_suite(void)
 {
     Suite *s = suite_create("wolfBoot");
     TCase *uart_flash = tcase_create("UART flash");
 
     tcase_add_test(uart_flash, test_ext_flash_read_timeout_returns_error);
+    tcase_add_test(uart_flash, test_ext_flash_erase_success);
+    tcase_add_test(uart_flash, test_ext_flash_erase_timeout_returns_error);
     tcase_set_timeout(uart_flash, 20);
     suite_add_tcase(s, uart_flash);