From 8f210ec145f7008e51d06dcf002ce1bad8f2efb8 Mon Sep 17 00:00:00 2001
From: Chris <firstof9@gmail.com>
Date: Wed, 13 Dec 2023 08:48:45 -0700
Subject: [PATCH] fix: handle unit of measure toggle option (#15)

* fix: handle unit of measure toggle option

* formatting
---
 custom_components/gasbuddy/__init__.py    | 32 +++++++++++++++++++++++
 custom_components/gasbuddy/config_flow.py |  2 +-
 tests/const.py                            |  1 -
 tests/test_init.py                        |  4 +--
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/custom_components/gasbuddy/__init__.py b/custom_components/gasbuddy/__init__.py
index 1cded89..07809bf 100644
--- a/custom_components/gasbuddy/__init__.py
+++ b/custom_components/gasbuddy/__init__.py
@@ -17,6 +17,7 @@
 from .const import (
     CONF_INTERVAL,
     CONF_STATION_ID,
+    CONF_UOM,
     COORDINATOR,
     DOMAIN,
     ISSUE_URL,
@@ -43,6 +44,14 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
         ISSUE_URL,
     )
 
+    # Some sanity checks
+    updated_config = config_entry.data.copy()
+    if CONF_UOM not in config_entry.data.keys():
+        updated_config[CONF_UOM] = True
+
+    if updated_config != config_entry.data:
+        hass.config_entries.async_update_entry(config_entry, data=updated_config)
+
     config_entry.add_update_listener(update_listener)
     interval = config_entry.data.get(CONF_INTERVAL)
     coordinator = GasBuddyUpdateCoordinator(hass, interval, config_entry)
@@ -72,6 +81,7 @@ async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> Non
         return
 
     config_entry.data[CONF_INTERVAL] = config_entry.options[CONF_INTERVAL]
+    config_entry.data[CONF_UOM] = config_entry.options[CONF_UOM]
 
     hass.config_entries.async_update_entry(
         entry=config_entry,
@@ -81,6 +91,28 @@ async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> Non
     await hass.config_entries.async_reload(config_entry.entry_id)
 
 
+async def async_migrate_entry(hass, config_entry) -> bool:
+    """Migrate an old config entry."""
+    version = config_entry.version
+
+    # 1 -> 2: Migrate format
+    if version == 1:
+        _LOGGER.debug("Migrating from version %s", version)
+        updated_config = config_entry.data.copy()
+
+        # Add default unit of measure setting if missing
+        if CONF_UOM not in updated_config.keys():
+            updated_config[CONF_UOM] = True
+
+        if updated_config != config_entry.data:
+            hass.config_entries.async_update_entry(config_entry, data=updated_config)
+
+        config_entry.version = 4
+        _LOGGER.debug("Migration to version %s complete", config_entry.version)
+
+    return True
+
+
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Handle removal of an entry."""
     _LOGGER.debug("Attempting to unload entities from the %s integration", DOMAIN)
diff --git a/custom_components/gasbuddy/config_flow.py b/custom_components/gasbuddy/config_flow.py
index bdd0541..b189328 100644
--- a/custom_components/gasbuddy/config_flow.py
+++ b/custom_components/gasbuddy/config_flow.py
@@ -168,7 +168,7 @@ def _get_default(key: str, fallback_default: Any = None) -> Any | None:
 class GasBuddyFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
     """Config flow for GasBuddy."""
 
-    VERSION = 1
+    VERSION = 2
     CONNECTION_CLASS = config_entries.CONN_CLASS_CLOUD_POLL
 
     def __init__(self):
diff --git a/tests/const.py b/tests/const.py
index 632a91a..03f397b 100644
--- a/tests/const.py
+++ b/tests/const.py
@@ -11,7 +11,6 @@
     CONF_NAME: "Gas Station",
     CONF_INTERVAL: 3600,
     CONF_STATION_ID: 208656,
-    CONF_UOM: True,
 }
 
 STATION_LIST = {
diff --git a/tests/test_init.py b/tests/test_init.py
index 763242f..b386afe 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -13,9 +13,7 @@
 async def test_setup_and_unload_entry(hass, mock_gasbuddy):
     """Test setup_entry."""
     entry = MockConfigEntry(
-        domain=DOMAIN,
-        title="gas_station",
-        data=CONFIG_DATA,
+        domain=DOMAIN, title="gas_station", data=CONFIG_DATA, version=1
     )
 
     entry.add_to_hass(hass)