Add restoring default effects

BazaJayGee66 · Feb 11, 2021 · 6b5c01b · 6b5c01b
1 parent 5b5c8bb
commit 6b5c01b
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 6 deletions.
diff --git a/custom_components/cololight/config_flow.py b/custom_components/cololight/config_flow.py
@@ -11,16 +11,16 @@
 
 from . import DOMAIN
 
-default_effects = list(PyCololight.DEFAULT_EFFECTS.keys())
+DEFAULT_EFFECTS = list(PyCololight.DEFAULT_EFFECTS.keys())
 
 DATA_SCHEMA = vol.Schema(
     {
         vol.Required(CONF_HOST): str,
         vol.Optional(CONF_NAME): str,
         vol.Optional(
             "default_effects",
-            default=default_effects,
-        ): cv.multi_select(default_effects),
+            default=DEFAULT_EFFECTS,
+        ): cv.multi_select(DEFAULT_EFFECTS),
     }
 )
 
@@ -80,6 +80,13 @@ def _get_effects(self):
         cololight = self.hass.data["cololight"][self.config_entry.entry_id]
         return dict(zip(cololight.effects, cololight.effects))
 
+    def _get_removed_effects(self):
+        cololight = self.hass.data["cololight"][self.config_entry.entry_id]
+        effects = cololight.effects
+        default_effects = DEFAULT_EFFECTS
+        removed_effects = list(set(default_effects) - set(effects))
+        return dict(zip(removed_effects, removed_effects))
+
     async def _is_valid(self, user_input):
         if not 1 <= user_input["cycle_speed"] <= 32:
             self._errors["cycle_speed"] = "invalid_cycle_speed"
@@ -95,12 +102,14 @@ async def async_step_init(self, user_input=None):
                 return await self.async_step_options_add_custom_effect()
             elif user_input["select"] == "Delete":
                 return await self.async_step_options_delete_effect()
+            elif user_input["select"] == "Restore":
+                return await self.async_step_options_restore_effect()
 
         options = {
             vol.Optional(
                 "select",
                 default=self.config_entry.options.get("select", "Create"),
-            ): vol.In(["Create", "Delete"]),
+            ): vol.In(["Create", "Delete", "Restore"]),
         }
 
         return self.async_show_form(step_id="init", data_schema=vol.Schema(options))
@@ -166,3 +175,24 @@ async def async_step_options_delete_effect(self, user_input=None):
         return self.async_show_form(
             step_id="options_delete_effect", data_schema=vol.Schema(options)
         )
+
+    async def async_step_options_restore_effect(self, user_input=None):
+        if user_input is not None:
+            for effect in user_input[CONF_NAME]:
+                self.config_entry.data["default_effects"].append(effect)
+
+            self.options["restored_effects"] = self.config_entry.data["default_effects"]
+
+            return self.async_create_entry(title="", data=self.options)
+
+        effects = self._get_removed_effects()
+        options = {
+            vol.Required(
+                CONF_NAME,
+                default=self.config_entry.options.get(CONF_NAME),
+            ): cv.multi_select(effects),
+        }
+
+        return self.async_show_form(
+            step_id="options_restore_effect", data_schema=vol.Schema(options)
+        )
diff --git a/custom_components/cololight/light.py b/custom_components/cololight/light.py
@@ -91,7 +91,7 @@ async def async_setup_entry(hass, entry, async_add_entities):
 
     if entry.options:
         for effect_name, effect_options in entry.options.items():
-            if effect_name != "default_effects":
+            if effect_name not in ["default_effects", "restored_effects"]:
                 try:
                     cololight_light.add_custom_effect(
                         effect_name,

diff --git a/custom_components/cololight/translations/en.json b/custom_components/cololight/translations/en.json
@@ -20,7 +20,7 @@
         "data": {
           "select": "Select action"
         },
-        "description": "Create/Delete Cololight effects"
+        "description": "Create/Delete/Restore Cololight effects"
       },
       "options_add_custom_effect": {
         "data": {
@@ -36,6 +36,12 @@
           "name": "Select effect/s to delete"
         },
         "description": "Delete Cololight effects"
+      },
+      "options_restore_effect": {
+        "data": {
+          "name": "Select effect/s to restore"
+        },
+        "description": "Restore default Cololight effects"
       }
     },
     "error": {

diff --git a/custom_components/tests/test_config_flow.py b/custom_components/tests/test_config_flow.py
@@ -33,6 +33,15 @@
     ],
 }
 
+DEMO_USER_INPUT_3 = {
+    "name": NAME,
+    "host": HOST,
+    "default_effects": [
+        "80s Club",
+        "Cherry Blossom",
+    ],
+}
+
 
 async def test_form(hass):
     """Test config entry configured successfully."""
@@ -290,6 +299,38 @@ async def test_options_deleting_default_effect(hass):
     assert hass.data["cololight"][entry["result"].entry_id].effects == ["80s Club"]
 
 
+async def test_options_restoring_default_effect(hass):
+    """Test options for restoring effect"""
+    flow = await hass.config_entries.flow.async_init(
+        cololight.DOMAIN, context={"source": "user"}
+    )
+    entry = await hass.config_entries.flow.async_configure(
+        flow["flow_id"], user_input=DEMO_USER_INPUT_3
+    )
+
+    await hass.async_block_till_done()
+
+    option_flow = await hass.config_entries.options.async_init(entry["result"].entry_id)
+
+    option_flow = await hass.config_entries.options.async_configure(
+        option_flow["flow_id"], user_input={"select": "Restore"}
+    )
+
+    await hass.config_entries.options.async_configure(
+        option_flow["flow_id"],
+        user_input={
+            "name": ["Unicorns"],
+        },
+    )
+
+    await hass.async_block_till_done()
+    assert hass.data["cololight"][entry["result"].entry_id].effects == [
+        "80s Club",
+        "Cherry Blossom",
+        "Unicorns",
+    ]
+
+
 @patch(
     "homeassistant.components.cololight.config_flow.CololightOptionsFlowHandler._get_color_schemes",
     return_value=["Mood | Green"],
@@ -403,3 +444,23 @@ async def test_end_to_end_with_options(hass):
 
     await hass.async_block_till_done()
     assert hass.data["cololight"][entry["result"].entry_id].effects == ["80s Club"]
+
+    option_flow = await hass.config_entries.options.async_init(entry["result"].entry_id)
+
+    option_flow = await hass.config_entries.options.async_configure(
+        option_flow["flow_id"], user_input={"select": "Restore"}
+    )
+
+    await hass.config_entries.options.async_configure(
+        option_flow["flow_id"],
+        user_input={
+            "name": ["Savasana", "Unicorns"],
+        },
+    )
+
+    await hass.async_block_till_done()
+    assert hass.data["cololight"][entry["result"].entry_id].effects == [
+        "80s Club",
+        "Savasana",
+        "Unicorns",
+    ]