From 07015ad527f81d81829190b31569a35a305676ce Mon Sep 17 00:00:00 2001
From: microDev <70126934+microDev1@users.noreply.github.com>
Date: Thu, 10 Dec 2020 16:42:45 +0530
Subject: [PATCH] add ability to switch boot partition

---
 ports/esp32s2/common-hal/ota/__init__.c | 6 ++++++
 shared-bindings/ota/__init__.c          | 7 +++++++
 shared-bindings/ota/__init__.h          | 1 +
 3 files changed, 14 insertions(+)

diff --git a/ports/esp32s2/common-hal/ota/__init__.c b/ports/esp32s2/common-hal/ota/__init__.c
index 97cc67b766..4ec1021317 100644
--- a/ports/esp32s2/common-hal/ota/__init__.c
+++ b/ports/esp32s2/common-hal/ota/__init__.c
@@ -141,3 +141,9 @@ void common_hal_ota_finish(void) {
         ota_reset();
     }
 }
+
+void common_hal_ota_switch(void) {
+    if (esp_ota_set_boot_partition(esp_ota_get_next_update_partition(NULL)) != ESP_OK) {
+        mp_raise_RuntimeError(translate("Unable to switch partition"));
+    }
+}
diff --git a/shared-bindings/ota/__init__.c b/shared-bindings/ota/__init__.c
index 84a46c83a2..46399dcb2e 100644
--- a/shared-bindings/ota/__init__.c
+++ b/shared-bindings/ota/__init__.c
@@ -30,6 +30,12 @@
 //|
 //| The `ota` module implements over-the-air update."""
 
+STATIC mp_obj_t ota_switch(void) {
+    common_hal_ota_switch();
+    return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(ota_switch_obj, ota_switch);
+
 STATIC mp_obj_t ota_finish(void) {
     common_hal_ota_finish();
     return mp_const_none;
@@ -60,6 +66,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_KW(ota_flash_obj, 1, ota_flash);
 
 STATIC const mp_rom_map_elem_t ota_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_ota) },
+    { MP_ROM_QSTR(MP_QSTR_switch), MP_ROM_PTR(&ota_switch_obj) },
     { MP_ROM_QSTR(MP_QSTR_finish), MP_ROM_PTR(&ota_finish_obj) },
     { MP_ROM_QSTR(MP_QSTR_flash), MP_ROM_PTR(&ota_flash_obj) },
 };
diff --git a/shared-bindings/ota/__init__.h b/shared-bindings/ota/__init__.h
index 9eec604e0b..82273f975a 100644
--- a/shared-bindings/ota/__init__.h
+++ b/shared-bindings/ota/__init__.h
@@ -29,6 +29,7 @@
 
 #include "py/runtime.h"
 
+extern void common_hal_ota_switch(void);
 extern void common_hal_ota_finish(void);
 extern void common_hal_ota_flash(const void *buf, const size_t len, const int32_t offset);